Initial snapshot: hand motion pipeline (Dyn-HaMR + dex-retargeting + SPIDER)
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
+11
@@ -0,0 +1,11 @@
|
||||
dataset_dir: ./data/amass/generative
|
||||
data:
|
||||
fps: 30
|
||||
clip_length: 128
|
||||
gender: neutral
|
||||
up: z
|
||||
|
||||
canonical: False
|
||||
unified_orientation: True
|
||||
single_motion: False
|
||||
normalize: True
|
||||
+36
@@ -0,0 +1,36 @@
|
||||
import os
|
||||
|
||||
import numpy as np
|
||||
import yaml
|
||||
|
||||
|
||||
class Struct:
|
||||
def __init__(self, **entries):
|
||||
self.__dict__.update(entries)
|
||||
|
||||
|
||||
class Arguments:
|
||||
def __init__(self, base_dir, config_path, filename='default.yaml'):
|
||||
with open(os.path.join(config_path, 'mano.yaml'), 'r') as f:
|
||||
smpl = yaml.safe_load(f)
|
||||
|
||||
self.smpl = Struct(**smpl)
|
||||
self.smpl.offsets['right'] = np.array(self.smpl.offsets['right'])
|
||||
self.smpl.offsets['left'] = np.array(self.smpl.offsets['left'])
|
||||
self.smpl.parents = np.array(self.smpl.parents).astype(np.int32)
|
||||
self.smpl.joint_num = 16 # len(self.smpl.joints_to_use)
|
||||
self.smpl.joints_to_use = np.array(self.smpl.joints_to_use)
|
||||
self.smpl.joints_to_use = np.arange(0, 63).reshape((-1, 3))[self.smpl.joints_to_use].reshape(-1)
|
||||
self.smpl.smpl_body_model = os.path.join(base_dir, '_DATA/data/mano')
|
||||
|
||||
self.filename = os.path.splitext(filename)[0]
|
||||
with open(os.path.join(config_path, filename), 'r') as f:
|
||||
config = yaml.safe_load(f)
|
||||
|
||||
for key, value in config.items():
|
||||
if isinstance(value, dict):
|
||||
setattr(self, key, Struct(**value))
|
||||
else:
|
||||
setattr(self, key, value)
|
||||
|
||||
self.json = config
|
||||
+639
@@ -0,0 +1,639 @@
|
||||
import os
|
||||
import glob
|
||||
import json
|
||||
import torch
|
||||
import joblib
|
||||
import random
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
from scipy import linalg
|
||||
from arguments import Arguments
|
||||
from nemf.losses import GeodesicLoss
|
||||
from nemf.generative import Architecture
|
||||
from nemf.fk import ForwardKinematicsLayer
|
||||
from scripts.train_gmm import MaxMixturePrior
|
||||
from human_body_prior.tools.omni_tools import copy2cpu as c2c
|
||||
from utils import FOOT_IDX, estimate_angular_velocity, estimate_linear_velocity
|
||||
from rotations import axis_angle_to_matrix, matrix_to_rotation_6d, rotation_6d_to_matrix
|
||||
|
||||
|
||||
def padding_prior_data(data):
|
||||
N, T, J = data['pos'].shape[:3]
|
||||
# pad inputs to match the training sequence length
|
||||
zero_padding = torch.zeros(N, args.data.clip_length - T, J, 3).to(prior.device)
|
||||
identity_padding = torch.zeros(N, args.data.clip_length - T, 3, 3).to(prior.device)
|
||||
identity_padding[:, :, 0, 0] = 1
|
||||
identity_padding[:, :, 1, 1] = 1
|
||||
identity_padding[:, :, 2, 2] = 1
|
||||
rotation_6d_padding = torch.zeros(N, args.data.clip_length - T, 6).to(prior.device)
|
||||
rotation_6d_padding[..., 0] = 1
|
||||
rotation_6d_padding[..., 4] = 1
|
||||
|
||||
pos = torch.cat((data['pos'], zero_padding), dim=1)
|
||||
velocity = torch.cat((data['velocity'], zero_padding), dim=1)
|
||||
angular = torch.cat((data['angular'], zero_padding), dim=1)
|
||||
root_orient = torch.cat((data['root_orient'], rotation_6d_padding), dim=1)
|
||||
root_vel = torch.cat((data['root_vel'], zero_padding[:, :, 0]), dim=1)
|
||||
rotmat_padding = rotation_6d_padding.unsqueeze(2).expand(-1, -1, J, -1)
|
||||
global_xform = torch.cat((data['global_xform'], rotmat_padding), dim=1)
|
||||
|
||||
data['pos'] = pos
|
||||
data['velocity'] = velocity
|
||||
data['global_xform'] = global_xform
|
||||
data['angular'] = angular
|
||||
data['root_orient'] = root_orient
|
||||
data['root_vel'] = root_vel
|
||||
|
||||
return data
|
||||
|
||||
|
||||
def load_data(files, padding=False, max_samples=400):
|
||||
poses = []
|
||||
trans = []
|
||||
visibility = []
|
||||
assert len(files) != 0, 'files not found'
|
||||
|
||||
max_samples = min(max_samples, len(files))
|
||||
|
||||
# import ipdb; ipdb.set_trace()
|
||||
|
||||
for f in files[:max_samples]:
|
||||
bdata = dict(np.load(f, allow_pickle=True))
|
||||
if 'r_wrist_orient' in bdata.keys() and 'pose_right_hand' in bdata.keys():
|
||||
|
||||
coef = len(bdata['r_wrist_orient']) // args.data.clip_length
|
||||
|
||||
for k_ in range(coef):
|
||||
root_orient = bdata['r_wrist_orient'][k_*args.data.clip_length:(k_+1)*args.data.clip_length]
|
||||
pose_body = bdata['pose_right_hand'][k_*args.data.clip_length:(k_+1)*args.data.clip_length]
|
||||
trans.append(bdata['r_wrist_location'][k_*args.data.clip_length:(k_+1)*args.data.clip_length])
|
||||
|
||||
poses.append(np.concatenate((root_orient, pose_body), axis=-1))
|
||||
|
||||
elif 'pose_right_hand' in bdata.keys():
|
||||
pass
|
||||
# poses.append(bdata['pose_right_hand'][:, :72])
|
||||
else:
|
||||
raise RuntimeError(f'missing pose parameters in the file: {f}')
|
||||
|
||||
|
||||
if 'visibility' in bdata.keys():
|
||||
visibility.append(bdata['visibility'])
|
||||
|
||||
# if 'poses' in bdata.keys():
|
||||
# poses.append(bdata['poses'][:, :72])
|
||||
# elif 'root_orient' in bdata.keys() and 'pose_body' in bdata.keys():
|
||||
# root_orient = bdata['root_orient']
|
||||
# pose_body = bdata['pose_body']
|
||||
# poses.append(np.concatenate((root_orient, pose_body), axis=-1))
|
||||
# else:
|
||||
# raise RuntimeError(f'missing pose parameters in the file: {f}')
|
||||
# trans.append(bdata['trans'])
|
||||
# if 'visibility' in bdata.keys():
|
||||
# visibility.append(bdata['visibility'])
|
||||
|
||||
trans = torch.from_numpy(np.asarray(trans, np.float32)).cuda() # global translation (N, T, 3)
|
||||
N, T = trans.shape[:2]
|
||||
poses = torch.from_numpy(np.asarray(poses, np.float32)).cuda()
|
||||
poses = poses.view(N, T, -1, 3) # axis-angle (N, T, J, 3)
|
||||
|
||||
|
||||
# poses = poses.view(N, T, 24, 3) # axis-angle (N, T, J, 3)
|
||||
|
||||
root_orient = poses[:, :, 0].clone()
|
||||
root_rotation = axis_angle_to_matrix(root_orient) # (N, T, 3, 3)
|
||||
poses[:, :, 0] = 0
|
||||
|
||||
rotmat = axis_angle_to_matrix(poses) # (N, T, J, 3, 3)
|
||||
angular = estimate_angular_velocity(rotmat.clone(), dt=1.0 / args.data.fps) # angular velocity of all the joints (N, T, J, 3)
|
||||
pos, global_xform = fk(rotmat.view(-1, 16, 3, 3))
|
||||
pos = pos.contiguous().view(N, T, 16, 3) # local joint positions (N, T, J, 3)
|
||||
global_xform = global_xform.view(N, T, 16, 4, 4)
|
||||
global_xform = global_xform[:, :, :, :3, :3] # global transformation matrix for each joint (N, T, J, 3, 3)
|
||||
velocity = estimate_linear_velocity(pos, dt=1.0 / args.data.fps) # linear velocity of all the joints (N, T, J, 3)
|
||||
|
||||
root_vel = estimate_linear_velocity(trans, dt=1.0 / args.data.fps) # linear velocity of the root joint (N, T, 3)
|
||||
|
||||
global_pos = torch.matmul(root_rotation.unsqueeze(2), pos.unsqueeze(-1)).squeeze(-1) # (N, T, J, 3)
|
||||
global_pos = global_pos + trans.unsqueeze(2)
|
||||
|
||||
data = {
|
||||
'pos': pos,
|
||||
'velocity': velocity,
|
||||
'global_xform': matrix_to_rotation_6d(global_xform),
|
||||
'angular': angular,
|
||||
'root_orient': matrix_to_rotation_6d(root_rotation),
|
||||
'root_vel': root_vel,
|
||||
'global_pos': global_pos,
|
||||
'rotmat': rotmat,
|
||||
'trans': trans
|
||||
}
|
||||
|
||||
if len(visibility) != 0:
|
||||
visibility = torch.from_numpy(np.asarray(visibility, bool)).cuda()
|
||||
padding = torch.zeros(visibility.shape[:2] + torch.Size([2]), dtype=torch.bool).cuda()
|
||||
visibility = torch.cat((visibility, padding), dim=-1) # (T, 24)
|
||||
print(f'visibility: {visibility.shape}')
|
||||
data['visibility'] = visibility
|
||||
|
||||
if T < args.data.clip_length:
|
||||
padding = True
|
||||
|
||||
if padding:
|
||||
if T == args.data.clip_length:
|
||||
data['pos'] = data['pos'].view(2 * N, T // 2, 24, 3)
|
||||
data['velocity'] = data['velocity'].view(2 * N, T // 2, 24, 3)
|
||||
data['global_xform'] = data['global_xform'].view(2 * N, T // 2, 24, 6)
|
||||
data['angular'] = data['angular'].view(2 * N, T // 2, 24, 3)
|
||||
data['root_orient'] = data['root_orient'].view(2 * N, T // 2, 6)
|
||||
data['root_vel'] = data['root_vel'].view(2 * N, T // 2, 3)
|
||||
data['global_pos'] = data['global_pos'].view(2 * N, T // 2, 24, 3)
|
||||
data['rotmat'] = data['rotmat'].view(2 * N, T // 2, 24, 3, 3)
|
||||
data['trans'] = data['trans'].view(2 * N, T // 2, 3)
|
||||
|
||||
padding_data = padding_prior_data(data)
|
||||
for key in padding_data.keys():
|
||||
data[key] = padding_data[key]
|
||||
|
||||
return data
|
||||
|
||||
|
||||
def calculate_feature_statistics(feature):
|
||||
feature = feature.cpu().numpy()
|
||||
mu = np.mean(feature, axis=0)
|
||||
sigma = np.cov(feature, rowvar=False)
|
||||
|
||||
return mu, sigma
|
||||
|
||||
|
||||
def calculate_frechet_distance(mu1, sigma1, mu2, sigma2, eps=1e-6):
|
||||
"""Numpy implementation of the Frechet Distance.
|
||||
The Frechet distance between two multivariate Gaussians X_1 ~ N(mu_1, C_1)
|
||||
and X_2 ~ N(mu_2, C_2) is
|
||||
d^2 = ||mu_1 - mu_2||^2 + Tr(C_1 + C_2 - 2*sqrt(C_1*C_2)).
|
||||
Stable version by Dougal J. Sutherland.
|
||||
Params:
|
||||
-- mu1 : Numpy array containing the activations of a layer of the
|
||||
inception net (like returned by the function 'get_predictions')
|
||||
for generated samples.
|
||||
-- mu2 : The sample mean over activations, precalculated on an
|
||||
representative data set.
|
||||
-- sigma1: The covariance matrix over activations for generated samples.
|
||||
-- sigma2: The covariance matrix over activations, precalculated on an
|
||||
representative data set.
|
||||
Returns:
|
||||
-- : The Frechet Distance.
|
||||
"""
|
||||
|
||||
mu1 = np.atleast_1d(mu1)
|
||||
mu2 = np.atleast_1d(mu2)
|
||||
|
||||
sigma1 = np.atleast_2d(sigma1)
|
||||
sigma2 = np.atleast_2d(sigma2)
|
||||
|
||||
assert mu1.shape == mu2.shape, \
|
||||
'Training and test mean vectors have different lengths'
|
||||
assert sigma1.shape == sigma2.shape, \
|
||||
'Training and test covariances have different dimensions'
|
||||
|
||||
diff = mu1 - mu2
|
||||
|
||||
# Product might be almost singular
|
||||
covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False)
|
||||
if not np.isfinite(covmean).all():
|
||||
msg = ('fid calculation produces singular product; '
|
||||
'adding %s to diagonal of cov estimates') % eps
|
||||
print(msg)
|
||||
offset = np.eye(sigma1.shape[0]) * eps
|
||||
covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset))
|
||||
|
||||
# Numerical error might give slight imaginary component
|
||||
if np.iscomplexobj(covmean):
|
||||
if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3):
|
||||
m = np.max(np.abs(covmean.imag))
|
||||
raise ValueError('Imaginary component {}'.format(m))
|
||||
covmean = covmean.real
|
||||
|
||||
tr_covmean = np.trace(covmean)
|
||||
|
||||
return (diff.dot(diff) + np.trace(sigma1) + np.trace(sigma2) - 2 * tr_covmean)
|
||||
|
||||
|
||||
def evaluate_fid(data, gt_data, num_motions=None, mask=None):
|
||||
if num_motions is None:
|
||||
index = list(range(data['pos'].shape[0]))
|
||||
index_gt = list(range(gt_data['pos'].shape[0]))
|
||||
else:
|
||||
# randomly pick [num_motions] motions from the data
|
||||
assert num_motions <= data['pos'].shape[0]
|
||||
assert num_motions <= gt_data['pos'].shape[0]
|
||||
|
||||
index = list(range(data['pos'].shape[0]))
|
||||
index = random.sample(index, num_motions)
|
||||
index_gt = list(range(gt_data['pos'].shape[0]))
|
||||
index_gt = random.sample(index_gt, num_motions)
|
||||
|
||||
data_selected = {}
|
||||
gt_data_selected = {}
|
||||
for key in data.keys():
|
||||
data_selected[key] = data[key][index].clone()
|
||||
gt_data_selected[key] = gt_data[key][index_gt].clone()
|
||||
|
||||
if mask is not None:
|
||||
# replace with ground truth
|
||||
for key in data.keys():
|
||||
data_selected[key][:, mask] = gt_data[key][:, mask].clone()
|
||||
|
||||
with torch.no_grad():
|
||||
prior.set_input(data_selected)
|
||||
f_l, _, _ = prior.encode_local()
|
||||
f_g, _, _ = prior.encode_global()
|
||||
feature = torch.cat((f_l, f_g), dim=-1)
|
||||
|
||||
prior.set_input(gt_data_selected)
|
||||
f_l_gt, _, _ = prior.encode_local()
|
||||
f_g_gt, _, _ = prior.encode_global()
|
||||
feature_gt = torch.cat((f_l_gt, f_g_gt), dim=-1)
|
||||
|
||||
statistics = calculate_feature_statistics(feature.view(feature.shape[0], -1))
|
||||
statistics_gt = calculate_feature_statistics(feature_gt.view(feature_gt.shape[0], -1))
|
||||
|
||||
fid = calculate_frechet_distance(statistics[0], statistics[1], statistics_gt[0], statistics_gt[1])
|
||||
print(f'FID: {fid:.4f}')
|
||||
|
||||
return fid
|
||||
|
||||
|
||||
def evaluate_diversity(data):
|
||||
index = list(range(data['pos'].shape[0]))
|
||||
random.shuffle(index)
|
||||
index1 = index[:data['pos'].shape[0] // 2]
|
||||
index2 = index[data['pos'].shape[0] // 2:]
|
||||
|
||||
with torch.no_grad():
|
||||
prior.set_input(data)
|
||||
f_l, _, _ = prior.encode_local()
|
||||
# f_g, _, _ = prior.encode_global()
|
||||
# feature = torch.cat((f_l, f_g), dim=-1)
|
||||
|
||||
feature = f_l.clone()
|
||||
|
||||
|
||||
diversity = 0
|
||||
for first_idx, second_idx in zip(index1, index2):
|
||||
diversity += torch.dist(feature[first_idx], feature[second_idx])
|
||||
diversity /= len(index1)
|
||||
|
||||
print(f'Diversity: {diversity.item():.4f}')
|
||||
|
||||
return diversity.item()
|
||||
|
||||
|
||||
def evaluate_recon_error(data, gt_data):
|
||||
criterion_geo = GeodesicLoss()
|
||||
|
||||
root_orient = rotation_6d_to_matrix(data['root_orient'])
|
||||
root_orient_gt = rotation_6d_to_matrix(gt_data['root_orient'])
|
||||
orientation_error = criterion_geo(root_orient.view(-1, 3, 3), root_orient_gt.view(-1, 3, 3))
|
||||
|
||||
rotation_error = criterion_geo(data['rotmat'][:, :, 1:].reshape(-1, 3, 3), gt_data['rotmat'][:, :, 1:].reshape(-1, 3, 3))
|
||||
position_error = torch.linalg.norm((data['pos'] - gt_data['pos']), dim=-1).mean()
|
||||
translation_error = torch.linalg.norm((data['trans'] - gt_data['trans']), dim=-1).mean()
|
||||
final_displacement_error = torch.linalg.norm(data['global_pos'][:, -1] - gt_data['global_pos'][:, -1], dim=-1).mean()
|
||||
|
||||
return {
|
||||
'rotation': c2c(rotation_error) * 180.0 / np.pi,
|
||||
'position': c2c(position_error) * 100.0,
|
||||
'orientation': c2c(orientation_error) * 180.0 / np.pi,
|
||||
'translation': c2c(translation_error) * 100.0,
|
||||
'final_displacement': c2c(final_displacement_error) * 100.0
|
||||
}
|
||||
|
||||
|
||||
def synthetic_trajectory(N_motions):
|
||||
|
||||
z_l = torch.normal(mean=0, std=1, size=(N_motions, args.local_prior.z_dim)).to(prior.device)
|
||||
z_g = torch.normal(mean=0, std=1, size=(N_motions, args.global_prior.z_dim)).to(prior.device)
|
||||
# z_g = torch.zeros((N_motions, args.global_prior.z_dim)).to(prior.device)
|
||||
|
||||
|
||||
with torch.no_grad():
|
||||
output = prior.decode(z_l, z_g, length=args.data.clip_length, step=1)
|
||||
|
||||
rotmat = output["rotmat"] # (N, T, J, 3, 3)
|
||||
angular = estimate_angular_velocity(rotmat.clone(), dt=1.0 / args.data.fps) # angular velocity of all the joints (N, T, J, 3)
|
||||
pos, global_xform = fk(rotmat.view(-1, 16, 3, 3))
|
||||
pos = pos.contiguous().view(N_motions, args.data.clip_length, 16, 3) # local joint positions (N, T, J, 3)
|
||||
global_xform = global_xform.view(N_motions, args.data.clip_length, 16, 4, 4)
|
||||
|
||||
global_xform = global_xform[:, :, :, :3, :3] # global transformation matrix for each joint (N, T, J, 3, 3)
|
||||
velocity = estimate_linear_velocity(pos, dt=1.0 / args.data.fps) # linear velocity of all the joints (N, T, J, 3)
|
||||
|
||||
root_rotation = rotmat[:, :, 0] # (N, T, 3, 3)
|
||||
global_pos = torch.matmul(root_rotation.unsqueeze(2), pos.unsqueeze(-1)).squeeze(-1) # (N, T, J, 3)
|
||||
|
||||
output_ = {'pos': pos,
|
||||
'velocity': velocity,
|
||||
'global_xform': matrix_to_rotation_6d(global_xform),
|
||||
'angular': angular,
|
||||
'root_orient': matrix_to_rotation_6d(root_rotation),
|
||||
'rotmat': rotmat}
|
||||
|
||||
return output_
|
||||
|
||||
|
||||
def ablation_study(path, subtask):
|
||||
|
||||
num_motions = 100
|
||||
print(os.path.join(path, 'TCDHands', '*.npz'))
|
||||
|
||||
test_data = load_data(sorted(glob.glob(os.path.join('./outputs/generative_copy/TCDHands', '*.npz')))) # test data TCDHands
|
||||
val_data = load_data(sorted(glob.glob(os.path.join('./outputs/generative_copy/GRAB', '*.npz')))) # test data TCDHands
|
||||
train_data = load_data(sorted(glob.glob(os.path.join('./outputs/generative_copy/SAMP', '*.npz')))) # test data TCDHands
|
||||
synthetic_data = synthetic_trajectory(N_motions=num_motions)
|
||||
|
||||
pca_aa = joblib.load("data/pca.pkl")
|
||||
pca_mean = torch.tensor(pca_aa.mean_).to("cuda")
|
||||
pca_cov = torch.tensor(pca_aa.get_covariance()).to("cuda")
|
||||
pca_sv = torch.tensor(pca_aa.singular_values_).to("cuda")
|
||||
pca_pc = torch.tensor(pca_aa.components_).to("cuda")
|
||||
|
||||
random_samples_low_dim = np.random.randn(num_motions, pca_aa.n_components_)
|
||||
pca_data = pca_aa.inverse_transform(random_samples_low_dim)
|
||||
|
||||
gmm_aa = MaxMixturePrior()
|
||||
gmm_data = gmm_aa.gmm.sample(num_motions)
|
||||
|
||||
identity = torch.eye(3)
|
||||
identity = identity.view(1, 1, 3, 3).repeat(num_motions, args.data.clip_length, 1, 1)
|
||||
|
||||
# convert to tensor
|
||||
gmm_data = torch.tensor(gmm_data[0].reshape(num_motions, args.data.clip_length, -1, 3))
|
||||
gmm_data = torch.cat((identity.unsqueeze(2), axis_angle_to_matrix(gmm_data)), dim=2)
|
||||
gmm_rot6d = matrix_to_rotation_6d(gmm_data).reshape(num_motions * args.data.clip_length, -1, 6)
|
||||
|
||||
# convert to tensor
|
||||
pca_data = torch.tensor(pca_data.reshape(num_motions, args.data.clip_length, -1, 3))
|
||||
pca_data = torch.cat((identity.unsqueeze(2), axis_angle_to_matrix(pca_data)), dim=2)
|
||||
pca_rot6d = matrix_to_rotation_6d(pca_data).reshape(num_motions * args.data.clip_length, -1, 6)
|
||||
|
||||
gmm_pos, gmm_xform = fk(gmm_rot6d.to(fk.device))
|
||||
pca_pos, pca_xform = fk(pca_rot6d.to(fk.device))
|
||||
|
||||
gmm_xform = matrix_to_rotation_6d(gmm_xform[:, :, :3, :3])
|
||||
pca_xform = matrix_to_rotation_6d(pca_xform[:, :, :3, :3])
|
||||
|
||||
gmm_velocity = estimate_linear_velocity(gmm_pos.unsqueeze(0), dt=1.0 / args.data.fps).squeeze(0)
|
||||
pca_velocity = estimate_linear_velocity(pca_pos.unsqueeze(0), dt=1.0 / args.data.fps).squeeze(0)
|
||||
|
||||
gmm_angular = estimate_angular_velocity(gmm_data.unsqueeze(0), dt=1.0 / args.data.fps).squeeze(0)
|
||||
pca_angular = estimate_angular_velocity(pca_data.unsqueeze(0), dt=1.0 / args.data.fps).squeeze(0)
|
||||
|
||||
gmm_dict = {'pos': gmm_pos.reshape(num_motions, args.data.clip_length, -1, 3),
|
||||
'global_xform': gmm_xform.reshape(num_motions, args.data.clip_length, -1, 6),
|
||||
'velocity': gmm_velocity.reshape(num_motions, args.data.clip_length, -1, 3),
|
||||
'angular': gmm_angular}
|
||||
|
||||
pca_dict = {'pos': pca_pos.reshape(num_motions, args.data.clip_length, -1, 3),
|
||||
'global_xform': pca_xform.reshape(num_motions, args.data.clip_length, -1, 6),
|
||||
'velocity': pca_velocity.reshape(num_motions, args.data.clip_length, -1, 3),
|
||||
'angular': pca_angular}
|
||||
|
||||
metrics = dict()
|
||||
|
||||
def f(data, num_motions, key):
|
||||
div = evaluate_diversity(data)
|
||||
if key in metrics.keys():
|
||||
pass
|
||||
else:
|
||||
metrics[key] = {'DIV': [div]}
|
||||
|
||||
print('Motion Synthesis:')
|
||||
div_test = f(test_data, num_motions, 'TEST_DATA (TDCHANDS)')
|
||||
div_val = f(val_data, num_motions, 'VAL_DATA (SAMP)')
|
||||
div_train = f(train_data, num_motions, 'TRAIN_DATA (GRAB)')
|
||||
div_synt = f(synthetic_data, num_motions, 'SYNTHETIC')
|
||||
div_gmm = f(gmm_dict, num_motions, 'GMM')
|
||||
div_pca = f(pca_dict, num_motions, 'PCA')
|
||||
|
||||
df = pd.DataFrame.from_dict(metrics)
|
||||
df.to_csv(os.path.join(path, f'diversity.csv'), index=False)
|
||||
|
||||
|
||||
|
||||
def comparison(path, subtask):
|
||||
gt_data = load_data(sorted(glob.glob(os.path.join(path, 'gt', '*.npz'))))
|
||||
|
||||
file_path = os.path.join(path, subtask)
|
||||
|
||||
metrics = dict()
|
||||
if subtask == 'recon':
|
||||
humor_data = load_data(sorted(glob.glob(os.path.join(file_path, 'humor', '*.npz'))))
|
||||
hm_vae_data = load_data(sorted(glob.glob(os.path.join(file_path, 'hm-vae', '*.npz'))))
|
||||
nemf_data = load_data(sorted(glob.glob(os.path.join(file_path, 'nemf', '*.npz'))))
|
||||
|
||||
metrics['HUMOR'] = evaluate_recon_error(humor_data, gt_data)
|
||||
metrics['HMVAE'] = evaluate_recon_error(hm_vae_data, gt_data)
|
||||
metrics['NEMF'] = evaluate_recon_error(nemf_data, gt_data)
|
||||
print('Motion Reconstruction:')
|
||||
print(json.dumps(metrics, sort_keys=True, indent=4))
|
||||
|
||||
df = pd.DataFrame.from_dict(metrics)
|
||||
df.to_csv(os.path.join(file_path, f'{subtask}.csv'))
|
||||
elif subtask == 'synth':
|
||||
humor_data = load_data(sorted(glob.glob(os.path.join(file_path, 'humor', '.*.npz'))), padding=True)
|
||||
hm_vae_data = load_data(sorted(glob.glob(os.path.join(file_path, 'hm-vae', '*.npz'))), padding=True)
|
||||
nemf_data = load_data(sorted(glob.glob(os.path.join(file_path, 'nemf', '*.npz'))), padding=True)
|
||||
|
||||
num_motions = 100
|
||||
|
||||
def f(data, num_motions, key):
|
||||
fid = evaluate_fid(data, gt_data, num_motions)
|
||||
div = evaluate_diversity(data)
|
||||
index = list(range(data['global_pos'].shape[0]))
|
||||
index = random.sample(index, num_motions)
|
||||
s = foot_skating(data['global_pos'][index])
|
||||
if key in metrics.keys():
|
||||
metrics[key]['FID'].append(fid)
|
||||
metrics[key]['DIV'].append(div)
|
||||
metrics[key]['FOOT'].append(s)
|
||||
else:
|
||||
metrics[key] = {'FID': [fid], 'DIV': [div], 'FOOT': [s]}
|
||||
|
||||
print('Motion Synthesis:')
|
||||
for replication in range(20):
|
||||
print(f'==================== Replication {replication} ====================')
|
||||
f(humor_data, num_motions, 'HUMOR')
|
||||
f(hm_vae_data, num_motions, 'HMVAE')
|
||||
f(nemf_data, num_motions, 'NEMF')
|
||||
|
||||
for key in metrics.keys():
|
||||
df = pd.DataFrame.from_dict(metrics[key])
|
||||
df.to_csv(os.path.join(file_path, f'{subtask}_{key}.csv'), index=False)
|
||||
else:
|
||||
raise RuntimeError(f'unsupported evaluation task for comparison: {subtask}')
|
||||
|
||||
|
||||
def inbetween_benchmarks(path, keyframes=False):
|
||||
trans_lengths = [5, 10, 15, 20] if keyframes else [10, 20, 30]
|
||||
|
||||
gt_data = load_data(sorted(glob.glob(os.path.join(path, 'gt', '*.npz'))))
|
||||
s = foot_skating(gt_data['global_pos'])
|
||||
print(f'average foot skating of the real motions: {s:.4f}')
|
||||
|
||||
def f(data, key, mask):
|
||||
fid = evaluate_fid(data, gt_data, mask=mask)
|
||||
s = foot_skating(data['global_pos'], mask=mask)
|
||||
if key in statistics.keys():
|
||||
statistics[key]['FID'].append(fid)
|
||||
statistics[key]['FOOT'].append(s)
|
||||
else:
|
||||
statistics[key] = {'FID': [fid], 'FOOT': [s]}
|
||||
|
||||
file_path = os.path.join(path, 'keyframe') if keyframes else os.path.join(path, 'clip')
|
||||
statistics = {}
|
||||
for t in trans_lengths:
|
||||
print(f'==================== Length {t} ====================')
|
||||
if keyframes:
|
||||
mask = None # don't mask data for keyframe inbetweening
|
||||
else:
|
||||
end_frame = (args.data.clip_length - t) // 2
|
||||
mask = list(range(end_frame)) + list(range(end_frame + t, args.data.clip_length))
|
||||
print(f'mask: {mask}')
|
||||
|
||||
slerp_data = load_data(sorted(glob.glob(os.path.join(file_path, 'slerp', f'{t}_frames', '*.npz'))))
|
||||
f(slerp_data, 'SLERP', mask)
|
||||
|
||||
if not keyframes:
|
||||
inertia_data = load_data(sorted(glob.glob(os.path.join(file_path, 'inertialization', f'{t}_frames', '*.npz'))))
|
||||
f(inertia_data, 'INERTIALIZATION', mask)
|
||||
|
||||
rmi_data = load_data(sorted(glob.glob(os.path.join(file_path, 'rmi', f'{t}_frames', '*.npz'))))
|
||||
f(rmi_data, 'RMI', mask)
|
||||
|
||||
hm_vae_data = load_data(sorted(glob.glob(os.path.join(file_path, 'hm-vae', f'{t}_frames', '*.npz'))))
|
||||
f(hm_vae_data, 'HMVAE', mask)
|
||||
|
||||
nemf_data = load_data(sorted(glob.glob(os.path.join(file_path, 'nemf', f'{t}_frames', '*.npz'))))
|
||||
f(nemf_data, 'NEMF', mask)
|
||||
|
||||
for key in statistics.keys():
|
||||
df = pd.DataFrame.from_dict(statistics[key])
|
||||
df.to_csv(os.path.join(file_path, f'inbetween_{key}.csv'))
|
||||
|
||||
|
||||
def compute_joint_accel(joint_seq, fps):
|
||||
''' Magnitude of joint accelerations for joint_seq : B x T x J x 3 '''
|
||||
joint_accel = joint_seq[:, :-2] - (2 * joint_seq[:, 1:-1]) + joint_seq[:, 2:]
|
||||
DATA_h = 1.0 / fps
|
||||
joint_accel = joint_accel / ((DATA_h**2))
|
||||
joint_accel_mag = torch.norm(joint_accel, dim=-1)
|
||||
|
||||
return joint_accel, joint_accel_mag
|
||||
|
||||
|
||||
def evaluate_super_sampling(path):
|
||||
fps_30_data = load_data(sorted(glob.glob(os.path.join(path, '30fps/*.npz'))))
|
||||
fps_60_data = load_data(sorted(glob.glob(os.path.join(path, '60fps/*.npz'))))
|
||||
fps_120_data = load_data(sorted(glob.glob(os.path.join(path, '120fps/*.npz'))))
|
||||
fps_240_data = load_data(sorted(glob.glob(os.path.join(path, '240fps/*.npz'))))
|
||||
|
||||
smoothness = dict()
|
||||
joint_vel = estimate_linear_velocity(fps_30_data['global_pos'], dt=1.0 / 30.0)
|
||||
joint_vel_mag = torch.norm(joint_vel, dim=-1).mean()
|
||||
smoothness['30FPS'] = [c2c(joint_vel_mag * 100.0)]
|
||||
joint_vel = estimate_linear_velocity(fps_60_data['global_pos'], dt=1.0 / 60.0)
|
||||
joint_vel_mag = torch.norm(joint_vel, dim=-1).mean()
|
||||
smoothness['60FPS'] = [c2c(joint_vel_mag * 100.0)]
|
||||
joint_vel = estimate_linear_velocity(fps_120_data['global_pos'], dt=1.0 / 120.0)
|
||||
joint_vel_mag = torch.norm(joint_vel, dim=-1).mean()
|
||||
smoothness['120FPS'] = [c2c(joint_vel_mag * 100.0)]
|
||||
joint_vel = estimate_linear_velocity(fps_240_data['global_pos'], dt=1.0 / 240.0)
|
||||
joint_vel_mag = torch.norm(joint_vel, dim=-1).mean()
|
||||
smoothness['240FPS'] = [c2c(joint_vel_mag * 100.0)]
|
||||
|
||||
df = pd.DataFrame.from_dict(smoothness)
|
||||
df.to_csv(os.path.join(path, f'smoothness.csv'))
|
||||
|
||||
|
||||
def evaluate_smoothness_test(path):
|
||||
smoothness = dict()
|
||||
|
||||
for fps in [30, 60]:
|
||||
smoothness[f'{fps}FPS'] = dict()
|
||||
for L in range(1, 23, 2):
|
||||
data = load_data([os.path.join(path, f'L{L}', f'test_{fps}fps.npz')])
|
||||
# _, joint_accel_mag = compute_joint_accel(data['global_pos'], fps=fps)
|
||||
# smoothness[f'{fps}FPS'][L] = c2c(joint_accel_mag.mean() * 100.0)
|
||||
joint_vel = estimate_linear_velocity(data['global_pos'], dt=1.0 / fps)
|
||||
joint_vel_mag = torch.norm(joint_vel, dim=-1).mean()
|
||||
smoothness[f'{fps}FPS'][L] = c2c(joint_vel_mag.mean() * 100.0)
|
||||
|
||||
df = pd.DataFrame.from_dict(smoothness)
|
||||
df.to_csv(os.path.join(path, f'smoothness.csv'))
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
plt.rcParams['text.usetex'] = True
|
||||
|
||||
smoothness_30fps = np.array(list(smoothness['30FPS'].values()))
|
||||
smoothness_60fps = np.array(list(smoothness['60FPS'].values()))
|
||||
|
||||
plt.plot(np.arange(1, 23, 2), smoothness_30fps, marker='o', label='30fps')
|
||||
plt.plot(np.arange(1, 23, 2), smoothness_60fps, marker='o', label='60fps')
|
||||
plt.legend()
|
||||
plt.grid(linestyle='--')
|
||||
plt.xticks(np.arange(1, 23, 2))
|
||||
plt.xlabel(r'$L$')
|
||||
# plt.yscale('log')
|
||||
plt.ylabel('Mean Per-Joint Velocity')
|
||||
plt.savefig(os.path.join(path, 'smoothness.pdf'))
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
from argparse import ArgumentParser
|
||||
|
||||
parser = ArgumentParser()
|
||||
parser.add_argument('--config', type=str, default='evaluation.yaml', help='name of the configutation file')
|
||||
parser.add_argument('--task', type=str, help='name of application tasks', required=True)
|
||||
parser.add_argument('--load_path', type=str, help='path to load results for evaluation', required=True)
|
||||
|
||||
opt = parser.parse_args()
|
||||
args = Arguments('./configs', filename=opt.config)
|
||||
|
||||
torch.set_default_dtype(torch.float32)
|
||||
|
||||
torch.manual_seed(0)
|
||||
np.random.seed(0)
|
||||
random.seed(0)
|
||||
|
||||
ngpu = 1
|
||||
if args.multi_gpu is True:
|
||||
ngpu = torch.cuda.device_count()
|
||||
if ngpu == 1:
|
||||
args.multi_gpu = False
|
||||
print(f'Number of GPUs: {ngpu}')
|
||||
|
||||
prior = Architecture(args, ngpu)
|
||||
prior.load(optimal=True)
|
||||
prior.eval()
|
||||
fk = ForwardKinematicsLayer(args)
|
||||
|
||||
if opt.task == 'ablation_study':
|
||||
print('Evaluating ablation study ...')
|
||||
ablation_study(opt.load_path, subtask='synth')
|
||||
elif opt.task == 'comparison':
|
||||
print('Evaluating comparison with other methods ...')
|
||||
comparison(opt.load_path, subtask='recon')
|
||||
comparison(opt.load_path, subtask='synth')
|
||||
elif opt.task == 'inbetween_benchmark':
|
||||
print('Evaluating motion in-betweening ...')
|
||||
inbetween_benchmarks(opt.load_path, keyframes=False)
|
||||
inbetween_benchmarks(opt.load_path, keyframes=True)
|
||||
elif opt.task == 'super_sampling':
|
||||
print('Evaluating super sampling ...')
|
||||
evaluate_super_sampling(opt.load_path)
|
||||
elif opt.task == 'smoothness_test':
|
||||
print('Evaluating smoothness testing ...')
|
||||
evaluate_smoothness_test(opt.load_path)
|
||||
else:
|
||||
raise NotImplementedError(f'unknown evaluation: {opt.task}')
|
||||
+1731
File diff suppressed because it is too large
Load Diff
+1119
File diff suppressed because it is too large
Load Diff
+150
@@ -0,0 +1,150 @@
|
||||
batch_size: 1
|
||||
bvh_viz: false
|
||||
|
||||
data:
|
||||
clip_length: 128
|
||||
fps: 30
|
||||
gender: neutral
|
||||
normalize:
|
||||
- pos
|
||||
- velocity
|
||||
- global_xform
|
||||
- angular
|
||||
- height
|
||||
- root_orient
|
||||
- root_vel
|
||||
root_transform: true
|
||||
up: z
|
||||
|
||||
root: None
|
||||
dataset_dir: _DATA/hmp_model
|
||||
dtw_loss: true
|
||||
geodesic_loss: true
|
||||
|
||||
global_prior:
|
||||
activation: tanh
|
||||
in_channels: 6
|
||||
kernel_size: 4
|
||||
temporal_scale: 8
|
||||
z_dim: 256
|
||||
|
||||
initialization: true
|
||||
is_train: false
|
||||
keypoint_source: mmpose
|
||||
|
||||
l1_loss: true
|
||||
lambda_angle: 1
|
||||
lambda_dtw: 0.5
|
||||
lambda_kl: 0.0001
|
||||
|
||||
local_prior:
|
||||
activation: tanh
|
||||
channel_base: 15
|
||||
extra_conv: 0
|
||||
kernel_size: 4
|
||||
num_layers: 4
|
||||
padding_mode: reflect
|
||||
skeleton_dist: 2
|
||||
skeleton_pool: mean
|
||||
temporal_scale: 8
|
||||
upsampling: linear
|
||||
use_residual_blocks: true
|
||||
z_dim: 1024
|
||||
|
||||
log: false
|
||||
multi_gpu: false
|
||||
|
||||
nemf:
|
||||
bandwidth: 7
|
||||
global_output: 6
|
||||
global_z: 256
|
||||
hidden_neuron: 1024
|
||||
local_output: 144
|
||||
local_z: 1024
|
||||
norm_layer: true
|
||||
siren: false
|
||||
skip_connection: true
|
||||
nsubject: 1
|
||||
num_workers: 0
|
||||
opt_betas: false
|
||||
output_trans: false
|
||||
overlap_len: 16
|
||||
pretrained_gmp: false
|
||||
save_dir: _DATA/hmp_model
|
||||
|
||||
scheduler:
|
||||
gamma: 0.7
|
||||
step_size: 200
|
||||
|
||||
stg1:
|
||||
betas_prior: 10
|
||||
keypoint_source: mmpose
|
||||
lambda_batch_cs: 0
|
||||
lambda_init_z_prior: 0
|
||||
lambda_j3d_smooth: 0
|
||||
lambda_motion_prior: 0
|
||||
lambda_pose_prior: 0
|
||||
lambda_orient: 0
|
||||
lambda_orient_smooth: 1
|
||||
lambda_pos: 0
|
||||
lambda_reproj: 0.05
|
||||
lambda_bio: 1
|
||||
lambda_consistency: 0
|
||||
lambda_rot: 0
|
||||
lambda_rot_smooth: 0
|
||||
lambda_trans: 10
|
||||
lambda_trans_smooth: 10
|
||||
lr: 0.05
|
||||
niters: 400
|
||||
opt_params:
|
||||
- betas
|
||||
- trans
|
||||
- root_orient
|
||||
|
||||
stg2:
|
||||
betas_prior: 10
|
||||
keypoint_source: mmpose
|
||||
lambda_batch_cs: 0
|
||||
lambda_init_z_prior: 0
|
||||
lambda_j3d_smooth: 0
|
||||
lambda_motion_prior: 200
|
||||
lambda_pose_prior: 0
|
||||
lambda_orient: 0
|
||||
lambda_orient_smooth: 1
|
||||
lambda_pos: 0
|
||||
lambda_reproj: 0.05
|
||||
lambda_bio: 1
|
||||
lambda_consistency: 0
|
||||
lambda_rot: 0
|
||||
lambda_rot_smooth: 0
|
||||
lambda_trans: 10
|
||||
lambda_trans_smooth: 10
|
||||
lr: 0.05
|
||||
niters: 400
|
||||
opt_params:
|
||||
- betas
|
||||
- trans
|
||||
- root_orient
|
||||
- z_l
|
||||
|
||||
stg3:
|
||||
betas_prior: 0
|
||||
lambda_batch_cs: 0
|
||||
lambda_init_z_prior: 0
|
||||
lambda_j3d_smooth: 0
|
||||
lambda_motion_prior: 0
|
||||
lambda_orient: 0
|
||||
lambda_orient_smooth: 0
|
||||
lambda_pos: 0
|
||||
lambda_reproj: 0
|
||||
lambda_bio: 0
|
||||
lambda_rot: 0
|
||||
lambda_rot_smooth: 0
|
||||
lambda_trans: 0
|
||||
lambda_trans_smooth: 0
|
||||
lr: 0.0
|
||||
niters: 0
|
||||
opt_params:
|
||||
- root_orient
|
||||
- z_l
|
||||
verbose: false
|
||||
@@ -0,0 +1,675 @@
|
||||
import operator
|
||||
|
||||
import numpy as np
|
||||
import numpy.core.umath as ut # (for numpy 1.25)
|
||||
# import numpy.core.umath_tests as ut
|
||||
|
||||
import AnimationStructure
|
||||
from Quaternions import Quaternions
|
||||
|
||||
|
||||
class Animation:
|
||||
"""
|
||||
Animation is a numpy-like wrapper for animation data
|
||||
|
||||
Animation data consists of several arrays consisting
|
||||
of F frames and J joints.
|
||||
|
||||
The animation is specified by
|
||||
|
||||
rotations : (F, J) Quaternions | Joint Rotations
|
||||
positions : (F, J, 3) ndarray | Joint Positions
|
||||
|
||||
The base pose is specified by
|
||||
|
||||
orients : (J) Quaternions | Joint Orientations
|
||||
offsets : (J, 3) ndarray | Joint Offsets
|
||||
|
||||
And the skeletal structure is specified by
|
||||
|
||||
parents : (J) ndarray | Joint Parents
|
||||
"""
|
||||
|
||||
def __init__(self, rotations, positions, orients, offsets, parents):
|
||||
|
||||
self.rotations = rotations
|
||||
self.positions = positions
|
||||
self.orients = orients
|
||||
self.offsets = offsets
|
||||
self.parents = parents
|
||||
|
||||
def __op__(self, op, other):
|
||||
return Animation(
|
||||
op(self.rotations, other.rotations),
|
||||
op(self.positions, other.positions),
|
||||
op(self.orients, other.orients),
|
||||
op(self.offsets, other.offsets),
|
||||
op(self.parents, other.parents))
|
||||
|
||||
def __iop__(self, op, other):
|
||||
self.rotations = op(self.roations, other.rotations)
|
||||
self.positions = op(self.roations, other.positions)
|
||||
self.orients = op(self.orients, other.orients)
|
||||
self.offsets = op(self.offsets, other.offsets)
|
||||
self.parents = op(self.parents, other.parents)
|
||||
return self
|
||||
|
||||
def __sop__(self, op):
|
||||
return Animation(
|
||||
op(self.rotations),
|
||||
op(self.positions),
|
||||
op(self.orients),
|
||||
op(self.offsets),
|
||||
op(self.parents))
|
||||
|
||||
def __add__(self, other): return self.__op__(operator.add, other)
|
||||
def __sub__(self, other): return self.__op__(operator.sub, other)
|
||||
def __mul__(self, other): return self.__op__(operator.mul, other)
|
||||
def __div__(self, other): return self.__op__(operator.div, other)
|
||||
|
||||
def __abs__(self): return self.__sop__(operator.abs)
|
||||
def __neg__(self): return self.__sop__(operator.neg)
|
||||
|
||||
def __iadd__(self, other): return self.__iop__(operator.iadd, other)
|
||||
def __isub__(self, other): return self.__iop__(operator.isub, other)
|
||||
def __imul__(self, other): return self.__iop__(operator.imul, other)
|
||||
def __idiv__(self, other): return self.__iop__(operator.idiv, other)
|
||||
|
||||
def __len__(self): return len(self.rotations)
|
||||
|
||||
def __getitem__(self, k):
|
||||
if isinstance(k, tuple):
|
||||
return Animation(
|
||||
self.rotations[k],
|
||||
self.positions[k],
|
||||
self.orients[k[1:]],
|
||||
self.offsets[k[1:]],
|
||||
self.parents[k[1:]])
|
||||
else:
|
||||
return Animation(
|
||||
self.rotations[k],
|
||||
self.positions[k],
|
||||
self.orients,
|
||||
self.offsets,
|
||||
self.parents)
|
||||
|
||||
def __setitem__(self, k, v):
|
||||
if isinstance(k, tuple):
|
||||
self.rotations.__setitem__(k, v.rotations)
|
||||
self.positions.__setitem__(k, v.positions)
|
||||
self.orients.__setitem__(k[1:], v.orients)
|
||||
self.offsets.__setitem__(k[1:], v.offsets)
|
||||
self.parents.__setitem__(k[1:], v.parents)
|
||||
else:
|
||||
self.rotations.__setitem__(k, v.rotations)
|
||||
self.positions.__setitem__(k, v.positions)
|
||||
self.orients.__setitem__(k, v.orients)
|
||||
self.offsets.__setitem__(k, v.offsets)
|
||||
self.parents.__setitem__(k, v.parents)
|
||||
|
||||
@property
|
||||
def shape(self): return (self.rotations.shape[0], self.rotations.shape[1])
|
||||
|
||||
def copy(self): return Animation(
|
||||
self.rotations.copy(), self.positions.copy(),
|
||||
self.orients.copy(), self.offsets.copy(),
|
||||
self.parents.copy())
|
||||
|
||||
def repeat(self, *args, **kw):
|
||||
return Animation(
|
||||
self.rotations.repeat(*args, **kw),
|
||||
self.positions.repeat(*args, **kw),
|
||||
self.orients, self.offsets, self.parents)
|
||||
|
||||
def ravel(self):
|
||||
return np.hstack([
|
||||
self.rotations.log().ravel(),
|
||||
self.positions.ravel(),
|
||||
self.orients.log().ravel(),
|
||||
self.offsets.ravel()])
|
||||
|
||||
@classmethod
|
||||
def unravel(clas, anim, shape, parents):
|
||||
nf, nj = shape
|
||||
rotations = anim[nf*nj*0:nf*nj*3]
|
||||
positions = anim[nf*nj*3:nf*nj*6]
|
||||
orients = anim[nf*nj*6+nj*0:nf*nj*6+nj*3]
|
||||
offsets = anim[nf*nj*6+nj*3:nf*nj*6+nj*6]
|
||||
return cls(
|
||||
Quaternions.exp(rotations), positions,
|
||||
Quaternions.exp(orients), offsets,
|
||||
parents.copy())
|
||||
|
||||
|
||||
""" Maya Interaction """
|
||||
|
||||
|
||||
def load_to_maya(anim, names=None, radius=0.5):
|
||||
"""
|
||||
Load Animation Object into Maya as Joint Skeleton
|
||||
loads each frame as a new keyfame in maya.
|
||||
|
||||
If the animation is too slow or too fast perhaps
|
||||
the framerate needs adjusting before being loaded
|
||||
such that it matches the maya scene framerate.
|
||||
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
anim : Animation
|
||||
Animation to load into Scene
|
||||
|
||||
names : [str]
|
||||
Optional list of Joint names for Skeleton
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
List of Maya Joint Nodes loaded into scene
|
||||
"""
|
||||
|
||||
import pymel.core as pm
|
||||
|
||||
joints = []
|
||||
frames = range(1, len(anim)+1)
|
||||
|
||||
if names is None:
|
||||
names = ["joint_" + str(i) for i in range(len(anim.parents))]
|
||||
|
||||
for i, offset, orient, parent, name in zip(range(len(anim.offsets)), anim.offsets, anim.orients, anim.parents, names):
|
||||
|
||||
if parent < 0:
|
||||
pm.select(d=True)
|
||||
else:
|
||||
pm.select(joints[parent])
|
||||
|
||||
joint = pm.joint(n=name, p=offset, relative=True, radius=radius)
|
||||
joint.setOrientation([orient[1], orient[2], orient[3], orient[0]])
|
||||
|
||||
curvex = pm.nodetypes.AnimCurveTA(n=name + "_rotateX")
|
||||
curvey = pm.nodetypes.AnimCurveTA(n=name + "_rotateY")
|
||||
curvez = pm.nodetypes.AnimCurveTA(n=name + "_rotateZ")
|
||||
|
||||
jrotations = (-Quaternions(orient[np.newaxis]) * anim.rotations[:, i]).euler()
|
||||
curvex.addKeys(frames, jrotations[:, 0])
|
||||
curvey.addKeys(frames, jrotations[:, 1])
|
||||
curvez.addKeys(frames, jrotations[:, 2])
|
||||
|
||||
pm.connectAttr(curvex.output, joint.rotateX)
|
||||
pm.connectAttr(curvey.output, joint.rotateY)
|
||||
pm.connectAttr(curvez.output, joint.rotateZ)
|
||||
|
||||
offsetx = pm.nodetypes.AnimCurveTU(n=name + "_translateX")
|
||||
offsety = pm.nodetypes.AnimCurveTU(n=name + "_translateY")
|
||||
offsetz = pm.nodetypes.AnimCurveTU(n=name + "_translateZ")
|
||||
|
||||
offsetx.addKeys(frames, anim.positions[:, i, 0])
|
||||
offsety.addKeys(frames, anim.positions[:, i, 1])
|
||||
offsetz.addKeys(frames, anim.positions[:, i, 2])
|
||||
|
||||
pm.connectAttr(offsetx.output, joint.translateX)
|
||||
pm.connectAttr(offsety.output, joint.translateY)
|
||||
pm.connectAttr(offsetz.output, joint.translateZ)
|
||||
|
||||
joints.append(joint)
|
||||
|
||||
return joints
|
||||
|
||||
|
||||
def load_from_maya(root, start, end):
|
||||
"""
|
||||
Load Animation Object from Maya Joint Skeleton
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
root : PyNode
|
||||
Root Joint of Maya Skeleton
|
||||
|
||||
start, end : int, int
|
||||
Start and End frame index of Maya Animation
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
animation : Animation
|
||||
Loaded animation from maya
|
||||
|
||||
names : [str]
|
||||
Joint names from maya
|
||||
"""
|
||||
|
||||
import pymel.core as pm
|
||||
|
||||
original_time = pm.currentTime(q=True)
|
||||
pm.currentTime(start)
|
||||
|
||||
""" Build Structure """
|
||||
|
||||
names, parents = AnimationStructure.load_from_maya(root)
|
||||
descendants = AnimationStructure.descendants_list(parents)
|
||||
orients = Quaternions.id(len(names))
|
||||
offsets = np.array([pm.xform(j, q=True, translation=True) for j in names])
|
||||
|
||||
for j, name in enumerate(names):
|
||||
scale = pm.xform(pm.PyNode(name), q=True, scale=True, relative=True)
|
||||
if len(descendants[j]) == 0:
|
||||
continue
|
||||
offsets[descendants[j]] *= scale
|
||||
|
||||
""" Load Animation """
|
||||
|
||||
eulers = np.zeros((end-start, len(names), 3))
|
||||
positions = np.zeros((end-start, len(names), 3))
|
||||
rotations = Quaternions.id((end-start, len(names)))
|
||||
|
||||
for i in range(end-start):
|
||||
|
||||
pm.currentTime(start+i+1, u=True)
|
||||
|
||||
scales = {}
|
||||
|
||||
for j, name, parent in zip(range(len(names)), names, parents):
|
||||
|
||||
node = pm.PyNode(name)
|
||||
|
||||
if i == 0 and pm.hasAttr(node, 'jointOrient'):
|
||||
ort = node.getOrientation()
|
||||
orients[j] = Quaternions(np.array([ort[3], ort[0], ort[1], ort[2]]))
|
||||
|
||||
if pm.hasAttr(node, 'rotate'):
|
||||
eulers[i, j] = np.radians(pm.xform(node, q=True, rotation=True))
|
||||
if pm.hasAttr(node, 'translate'):
|
||||
positions[i, j] = pm.xform(node, q=True, translation=True)
|
||||
if pm.hasAttr(node, 'scale'):
|
||||
scales[j] = pm.xform(node, q=True, scale=True, relative=True)
|
||||
|
||||
for j in scales:
|
||||
if len(descendants[j]) == 0:
|
||||
continue
|
||||
positions[i, descendants[j]] *= scales[j]
|
||||
|
||||
positions[i, 0] = pm.xform(root, q=True, translation=True, worldSpace=True)
|
||||
|
||||
rotations = orients[np.newaxis] * Quaternions.from_euler(eulers, order='xyz', world=True)
|
||||
|
||||
""" Done """
|
||||
|
||||
pm.currentTime(original_time)
|
||||
|
||||
return Animation(rotations, positions, orients, offsets, parents), names
|
||||
|
||||
# local transformation matrices
|
||||
|
||||
|
||||
def transforms_local(anim):
|
||||
"""
|
||||
Computes Animation Local Transforms
|
||||
|
||||
As well as a number of other uses this can
|
||||
be used to compute global joint transforms,
|
||||
which in turn can be used to compete global
|
||||
joint positions
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
anim : Animation
|
||||
Input animation
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
transforms : (F, J, 4, 4) ndarray
|
||||
|
||||
For each frame F, joint local
|
||||
transforms for each joint J
|
||||
"""
|
||||
|
||||
transforms = anim.rotations.transforms()
|
||||
transforms = np.concatenate([transforms, np.zeros(transforms.shape[:2] + (3, 1))], axis=-1)
|
||||
transforms = np.concatenate([transforms, np.zeros(transforms.shape[:2] + (1, 4))], axis=-2)
|
||||
# the last column is filled with the joint positions!
|
||||
transforms[:, :, 0:3, 3] = anim.positions
|
||||
transforms[:, :, 3:4, 3] = 1.0
|
||||
return transforms
|
||||
|
||||
|
||||
def transforms_multiply(t0s, t1s):
|
||||
"""
|
||||
Transforms Multiply
|
||||
|
||||
Multiplies two arrays of animation transforms
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
t0s, t1s : (F, J, 4, 4) ndarray
|
||||
Two arrays of transforms
|
||||
for each frame F and each
|
||||
joint J
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
transforms : (F, J, 4, 4) ndarray
|
||||
Array of transforms for each
|
||||
frame F and joint J multiplied
|
||||
together
|
||||
"""
|
||||
|
||||
return ut.matrix_multiply(t0s, t1s)
|
||||
|
||||
|
||||
def transforms_inv(ts):
|
||||
fts = ts.reshape(-1, 4, 4)
|
||||
fts = np.array(list(map(lambda x: np.linalg.inv(x), fts)))
|
||||
return fts.reshape(ts.shape)
|
||||
|
||||
|
||||
def transforms_blank(anim):
|
||||
"""
|
||||
Blank Transforms
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
anim : Animation
|
||||
Input animation
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
transforms : (F, J, 4, 4) ndarray
|
||||
Array of identity transforms for
|
||||
each frame F and joint J
|
||||
"""
|
||||
|
||||
ts = np.zeros(anim.shape + (4, 4))
|
||||
ts[:, :, 0, 0] = 1.0
|
||||
ts[:, :, 1, 1] = 1.0
|
||||
ts[:, :, 2, 2] = 1.0
|
||||
ts[:, :, 3, 3] = 1.0
|
||||
return ts
|
||||
|
||||
# global transformation matrices
|
||||
|
||||
|
||||
def transforms_global(anim):
|
||||
"""
|
||||
Global Animation Transforms
|
||||
|
||||
This relies on joint ordering
|
||||
being incremental. That means a joint
|
||||
J1 must not be a ancestor of J0 if
|
||||
J0 appears before J1 in the joint
|
||||
ordering.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
anim : Animation
|
||||
Input animation
|
||||
|
||||
Returns
|
||||
------
|
||||
|
||||
transforms : (F, J, 4, 4) ndarray
|
||||
Array of global transforms for
|
||||
each frame F and joint J
|
||||
"""
|
||||
|
||||
joints = np.arange(anim.shape[1])
|
||||
parents = np.arange(anim.shape[1])
|
||||
locals = transforms_local(anim)
|
||||
globals = transforms_blank(anim)
|
||||
|
||||
globals[:, 0] = locals[:, 0]
|
||||
|
||||
for i in range(1, anim.shape[1]):
|
||||
globals[:, i] = transforms_multiply(globals[:, anim.parents[i]], locals[:, i])
|
||||
|
||||
return globals
|
||||
|
||||
# !!! useful!
|
||||
|
||||
|
||||
def positions_global(anim):
|
||||
"""
|
||||
Global Joint Positions
|
||||
|
||||
Given an animation compute the global joint
|
||||
positions at at every frame
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
anim : Animation
|
||||
Input animation
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
positions : (F, J, 3) ndarray
|
||||
Positions for every frame F
|
||||
and joint position J
|
||||
"""
|
||||
|
||||
# get the last column -- corresponding to the coordinates
|
||||
positions = transforms_global(anim)[:, :, :, 3]
|
||||
return positions[:, :, :3] / positions[:, :, 3, np.newaxis]
|
||||
|
||||
|
||||
""" Rotations """
|
||||
|
||||
|
||||
def rotations_global(anim):
|
||||
"""
|
||||
Global Animation Rotations
|
||||
|
||||
This relies on joint ordering
|
||||
being incremental. That means a joint
|
||||
J1 must not be a ancestor of J0 if
|
||||
J0 appears before J1 in the joint
|
||||
ordering.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
anim : Animation
|
||||
Input animation
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
points : (F, J) Quaternions
|
||||
global rotations for every frame F
|
||||
and joint J
|
||||
"""
|
||||
|
||||
joints = np.arange(anim.shape[1])
|
||||
parents = np.arange(anim.shape[1])
|
||||
locals = anim.rotations
|
||||
globals = Quaternions.id(anim.shape)
|
||||
|
||||
globals[:, 0] = locals[:, 0]
|
||||
|
||||
for i in range(1, anim.shape[1]):
|
||||
globals[:, i] = globals[:, anim.parents[i]] * locals[:, i]
|
||||
|
||||
return globals
|
||||
|
||||
|
||||
def rotations_parents_global(anim):
|
||||
rotations = rotations_global(anim)
|
||||
rotations = rotations[:, anim.parents]
|
||||
rotations[:, 0] = Quaternions.id(len(anim))
|
||||
return rotations
|
||||
|
||||
|
||||
def rotations_load_to_maya(rotations, positions, names=None):
|
||||
"""
|
||||
Load Rotations into Maya
|
||||
|
||||
Loads a Quaternions array into the scene
|
||||
via the representation of axis
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
rotations : (F, J) Quaternions
|
||||
array of rotations to load
|
||||
into the scene where
|
||||
F = number of frames
|
||||
J = number of joints
|
||||
|
||||
positions : (F, J, 3) ndarray
|
||||
array of positions to load
|
||||
rotation axis at where:
|
||||
F = number of frames
|
||||
J = number of joints
|
||||
|
||||
names : [str]
|
||||
List of joint names
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
maxies : Group
|
||||
Grouped Maya Node of all Axis nodes
|
||||
"""
|
||||
|
||||
import pymel.core as pm
|
||||
|
||||
if names is None:
|
||||
names = ["joint_" + str(i) for i in range(rotations.shape[1])]
|
||||
|
||||
maxis = []
|
||||
frames = range(1, len(positions)+1)
|
||||
for i, name in enumerate(names):
|
||||
|
||||
name = name + "_axis"
|
||||
axis = pm.group(
|
||||
pm.curve(p=[(0, 0, 0), (1, 0, 0)], d=1, n=name+'_axis_x'),
|
||||
pm.curve(p=[(0, 0, 0), (0, 1, 0)], d=1, n=name+'_axis_y'),
|
||||
pm.curve(p=[(0, 0, 0), (0, 0, 1)], d=1, n=name+'_axis_z'),
|
||||
n=name)
|
||||
|
||||
axis.rotatePivot.set((0, 0, 0))
|
||||
axis.scalePivot.set((0, 0, 0))
|
||||
axis.childAtIndex(0).overrideEnabled.set(1)
|
||||
axis.childAtIndex(0).overrideColor.set(13)
|
||||
axis.childAtIndex(1).overrideEnabled.set(1)
|
||||
axis.childAtIndex(1).overrideColor.set(14)
|
||||
axis.childAtIndex(2).overrideEnabled.set(1)
|
||||
axis.childAtIndex(2).overrideColor.set(15)
|
||||
|
||||
curvex = pm.nodetypes.AnimCurveTA(n=name + "_rotateX")
|
||||
curvey = pm.nodetypes.AnimCurveTA(n=name + "_rotateY")
|
||||
curvez = pm.nodetypes.AnimCurveTA(n=name + "_rotateZ")
|
||||
|
||||
arotations = rotations[:, i].euler()
|
||||
curvex.addKeys(frames, arotations[:, 0])
|
||||
curvey.addKeys(frames, arotations[:, 1])
|
||||
curvez.addKeys(frames, arotations[:, 2])
|
||||
|
||||
pm.connectAttr(curvex.output, axis.rotateX)
|
||||
pm.connectAttr(curvey.output, axis.rotateY)
|
||||
pm.connectAttr(curvez.output, axis.rotateZ)
|
||||
|
||||
offsetx = pm.nodetypes.AnimCurveTU(n=name + "_translateX")
|
||||
offsety = pm.nodetypes.AnimCurveTU(n=name + "_translateY")
|
||||
offsetz = pm.nodetypes.AnimCurveTU(n=name + "_translateZ")
|
||||
|
||||
offsetx.addKeys(frames, positions[:, i, 0])
|
||||
offsety.addKeys(frames, positions[:, i, 1])
|
||||
offsetz.addKeys(frames, positions[:, i, 2])
|
||||
|
||||
pm.connectAttr(offsetx.output, axis.translateX)
|
||||
pm.connectAttr(offsety.output, axis.translateY)
|
||||
pm.connectAttr(offsetz.output, axis.translateZ)
|
||||
|
||||
maxis.append(axis)
|
||||
|
||||
return pm.group(*maxis, n='RotationAnimation')
|
||||
|
||||
|
||||
""" Offsets & Orients """
|
||||
|
||||
|
||||
def orients_global(anim):
|
||||
|
||||
joints = np.arange(anim.shape[1])
|
||||
parents = np.arange(anim.shape[1])
|
||||
locals = anim.orients
|
||||
globals = Quaternions.id(anim.shape[1])
|
||||
|
||||
globals[:, 0] = locals[:, 0]
|
||||
|
||||
for i in range(1, anim.shape[1]):
|
||||
globals[:, i] = globals[:, anim.parents[i]] * locals[:, i]
|
||||
|
||||
return globals
|
||||
|
||||
|
||||
def offsets_transforms_local(anim):
|
||||
|
||||
transforms = anim.orients[np.newaxis].transforms()
|
||||
transforms = np.concatenate([transforms, np.zeros(transforms.shape[:2] + (3, 1))], axis=-1)
|
||||
transforms = np.concatenate([transforms, np.zeros(transforms.shape[:2] + (1, 4))], axis=-2)
|
||||
transforms[:, :, 0:3, 3] = anim.offsets[np.newaxis]
|
||||
transforms[:, :, 3:4, 3] = 1.0
|
||||
return transforms
|
||||
|
||||
|
||||
def offsets_transforms_global(anim):
|
||||
|
||||
joints = np.arange(anim.shape[1])
|
||||
parents = np.arange(anim.shape[1])
|
||||
locals = offsets_transforms_local(anim)
|
||||
globals = transforms_blank(anim)
|
||||
|
||||
globals[:, 0] = locals[:, 0]
|
||||
|
||||
for i in range(1, anim.shape[1]):
|
||||
globals[:, i] = transforms_multiply(globals[:, anim.parents[i]], locals[:, i])
|
||||
|
||||
return globals
|
||||
|
||||
|
||||
def offsets_global(anim):
|
||||
offsets = offsets_transforms_global(anim)[:, :, :, 3]
|
||||
return offsets[0, :, :3] / offsets[0, :, 3, np.newaxis]
|
||||
|
||||
|
||||
""" Lengths """
|
||||
|
||||
|
||||
def offset_lengths(anim):
|
||||
return np.sum(anim.offsets[1:]**2.0, axis=1)**0.5
|
||||
|
||||
|
||||
def position_lengths(anim):
|
||||
return np.sum(anim.positions[:, 1:]**2.0, axis=2)**0.5
|
||||
|
||||
|
||||
""" Skinning """
|
||||
|
||||
|
||||
def skin(anim, rest, weights, mesh, maxjoints=4):
|
||||
|
||||
full_transforms = transforms_multiply(
|
||||
transforms_global(anim),
|
||||
transforms_inv(transforms_global(rest[0:1])))
|
||||
|
||||
weightids = np.argsort(-weights, axis=1)[:, :maxjoints]
|
||||
weightvls = np.array(list(map(lambda w, i: w[i], weights, weightids)))
|
||||
weightvls = weightvls / weightvls.sum(axis=1)[..., np.newaxis]
|
||||
|
||||
verts = np.hstack([mesh, np.ones((len(mesh), 1))])
|
||||
verts = verts[np.newaxis, :, np.newaxis, :, np.newaxis]
|
||||
verts = transforms_multiply(full_transforms[:, weightids], verts)
|
||||
verts = (verts[:, :, :, :3] / verts[:, :, :, 3:4])[:, :, :, :, 0]
|
||||
|
||||
return np.sum(weightvls[np.newaxis, :, :, np.newaxis] * verts, axis=2)
|
||||
@@ -0,0 +1,451 @@
|
||||
import numpy as np
|
||||
import scipy.sparse as sparse
|
||||
import Animation as Animation
|
||||
|
||||
|
||||
""" Maya Functions """
|
||||
|
||||
def load_from_maya(root):
|
||||
"""
|
||||
Load joint parents and names from maya
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
root : PyNode
|
||||
Root Maya Node
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
(names, parents) : ([str], (J) ndarray)
|
||||
List of joint names and array
|
||||
of indices representing the parent
|
||||
joint for each joint J.
|
||||
|
||||
Joint index -1 is used to represent
|
||||
that there is no parent joint
|
||||
"""
|
||||
|
||||
import pymel.core as pm
|
||||
|
||||
names = []
|
||||
parents = []
|
||||
|
||||
def unload_joint(j, parents, par):
|
||||
|
||||
id = len(names)
|
||||
names.append(j)
|
||||
parents.append(par)
|
||||
|
||||
children = [c for c in j.getChildren() if
|
||||
isinstance(c, pm.nt.Transform) and
|
||||
not isinstance(c, pm.nt.Constraint) and
|
||||
not any(pm.listRelatives(c, s=True)) and
|
||||
(any(pm.listRelatives(c, ad=True, ap=False, type='joint')) or isinstance(c, pm.nt.Joint))]
|
||||
|
||||
map(lambda c: unload_joint(c, parents, id), children)
|
||||
|
||||
unload_joint(root, parents, -1)
|
||||
|
||||
return (names, parents)
|
||||
|
||||
|
||||
""" Family Functions """
|
||||
|
||||
def joints(parents):
|
||||
"""
|
||||
Parameters
|
||||
----------
|
||||
|
||||
parents : (J) ndarray
|
||||
parents array
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
joints : (J) ndarray
|
||||
Array of joint indices
|
||||
"""
|
||||
return np.arange(len(parents), dtype=int)
|
||||
|
||||
def joints_list(parents):
|
||||
"""
|
||||
Parameters
|
||||
----------
|
||||
|
||||
parents : (J) ndarray
|
||||
parents array
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
joints : [ndarray]
|
||||
List of arrays of joint idices for
|
||||
each joint
|
||||
"""
|
||||
return list(joints(parents)[:,np.newaxis])
|
||||
|
||||
def parents_list(parents):
|
||||
"""
|
||||
Parameters
|
||||
----------
|
||||
|
||||
parents : (J) ndarray
|
||||
parents array
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
parents : [ndarray]
|
||||
List of arrays of joint idices for
|
||||
the parents of each joint
|
||||
"""
|
||||
return list(parents[:,np.newaxis])
|
||||
|
||||
|
||||
def children_list(parents):
|
||||
"""
|
||||
Parameters
|
||||
----------
|
||||
|
||||
parents : (J) ndarray
|
||||
parents array
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
children : [ndarray]
|
||||
List of arrays of joint indices for
|
||||
the children of each joint
|
||||
"""
|
||||
|
||||
def joint_children(i):
|
||||
return [j for j, p in enumerate(parents) if p == i]
|
||||
|
||||
return list(map(lambda j: np.array(joint_children(j)), joints(parents)))
|
||||
|
||||
|
||||
def descendants_list(parents):
|
||||
"""
|
||||
Parameters
|
||||
----------
|
||||
|
||||
parents : (J) ndarray
|
||||
parents array
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
descendants : [ndarray]
|
||||
List of arrays of joint idices for
|
||||
the descendants of each joint
|
||||
"""
|
||||
|
||||
children = children_list(parents)
|
||||
|
||||
def joint_descendants(i):
|
||||
return sum([joint_descendants(j) for j in children[i]], list(children[i]))
|
||||
|
||||
return list(map(lambda j: np.array(joint_descendants(j)), joints(parents)))
|
||||
|
||||
|
||||
def ancestors_list(parents):
|
||||
"""
|
||||
Parameters
|
||||
----------
|
||||
|
||||
parents : (J) ndarray
|
||||
parents array
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
ancestors : [ndarray]
|
||||
List of arrays of joint idices for
|
||||
the ancestors of each joint
|
||||
"""
|
||||
|
||||
decendants = descendants_list(parents)
|
||||
|
||||
def joint_ancestors(i):
|
||||
return [j for j in joints(parents) if i in decendants[j]]
|
||||
|
||||
return list(map(lambda j: np.array(joint_ancestors(j)), joints(parents)))
|
||||
|
||||
|
||||
""" Mask Functions """
|
||||
|
||||
def mask(parents, filter):
|
||||
"""
|
||||
Constructs a Mask for a give filter
|
||||
|
||||
A mask is a (J, J) ndarray truth table for a given
|
||||
condition over J joints. For example there
|
||||
may be a mask specifying if a joint N is a
|
||||
child of another joint M.
|
||||
|
||||
This could be constructed into a mask using
|
||||
`m = mask(parents, children_list)` and the condition
|
||||
of childhood tested using `m[N, M]`.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
parents : (J) ndarray
|
||||
parents array
|
||||
|
||||
filter : (J) ndarray -> [ndarray]
|
||||
function that outputs a list of arrays
|
||||
of joint indices for some condition
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
mask : (N, N) ndarray
|
||||
boolean truth table of given condition
|
||||
"""
|
||||
m = np.zeros((len(parents), len(parents))).astype(bool)
|
||||
jnts = joints(parents)
|
||||
fltr = filter(parents)
|
||||
for i,f in enumerate(fltr): m[i,:] = np.any(jnts[:,np.newaxis] == f[np.newaxis,:], axis=1)
|
||||
return m
|
||||
|
||||
def joints_mask(parents): return np.eye(len(parents)).astype(bool)
|
||||
def children_mask(parents): return mask(parents, children_list)
|
||||
def parents_mask(parents): return mask(parents, parents_list)
|
||||
def descendants_mask(parents): return mask(parents, descendants_list)
|
||||
def ancestors_mask(parents): return mask(parents, ancestors_list)
|
||||
|
||||
""" Search Functions """
|
||||
|
||||
def joint_chain_ascend(parents, start, end):
|
||||
chain = []
|
||||
while start != end:
|
||||
chain.append(start)
|
||||
start = parents[start]
|
||||
chain.append(end)
|
||||
return np.array(chain, dtype=int)
|
||||
|
||||
|
||||
""" Constraints """
|
||||
|
||||
def constraints(anim, **kwargs):
|
||||
"""
|
||||
Constraint list for Animation
|
||||
|
||||
This constraint list can be used in the
|
||||
VerletParticle solver to constrain
|
||||
a animation global joint positions.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
anim : Animation
|
||||
Input animation
|
||||
|
||||
masses : (F, J) ndarray
|
||||
Optional list of masses
|
||||
for joints J across frames F
|
||||
defaults to weighting by
|
||||
vertical height
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
constraints : [(int, int, (F, J) ndarray, (F, J) ndarray, (F, J) ndarray)]
|
||||
A list of constraints in the format:
|
||||
(Joint1, Joint2, Masses1, Masses2, Lengths)
|
||||
|
||||
"""
|
||||
|
||||
masses = kwargs.pop('masses', None)
|
||||
|
||||
children = children_list(anim.parents)
|
||||
constraints = []
|
||||
|
||||
points_offsets = Animation.offsets_global(anim)
|
||||
points = Animation.positions_global(anim)
|
||||
|
||||
if masses is None:
|
||||
masses = 1.0 / (0.1 + np.absolute(points_offsets[:,1]))
|
||||
masses = masses[np.newaxis].repeat(len(anim), axis=0)
|
||||
|
||||
for j in xrange(anim.shape[1]):
|
||||
|
||||
""" Add constraints between all joints and their children """
|
||||
for c0 in children[j]:
|
||||
|
||||
dists = np.sum((points[:, c0] - points[:, j])**2.0, axis=1)**0.5
|
||||
constraints.append((c0, j, masses[:,c0], masses[:,j], dists))
|
||||
|
||||
""" Add constraints between all children of joint """
|
||||
for c1 in children[j]:
|
||||
if c0 == c1: continue
|
||||
|
||||
dists = np.sum((points[:, c0] - points[:, c1])**2.0, axis=1)**0.5
|
||||
constraints.append((c0, c1, masses[:,c0], masses[:,c1], dists))
|
||||
|
||||
return constraints
|
||||
|
||||
""" Graph Functions """
|
||||
|
||||
def graph(anim):
|
||||
"""
|
||||
Generates a weighted adjacency matrix
|
||||
using local joint distances along
|
||||
the skeletal structure.
|
||||
|
||||
Joints which are not connected
|
||||
are assigned the weight `0`.
|
||||
|
||||
Joints which actually have zero distance
|
||||
between them, but are still connected, are
|
||||
perturbed by some minimal amount.
|
||||
|
||||
The output of this routine can be used
|
||||
with the `scipy.sparse.csgraph`
|
||||
routines for graph analysis.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
anim : Animation
|
||||
input animation
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
graph : (N, N) ndarray
|
||||
weight adjacency matrix using
|
||||
local distances along the
|
||||
skeletal structure from joint
|
||||
N to joint M. If joints are not
|
||||
directly connected are assigned
|
||||
the weight `0`.
|
||||
"""
|
||||
|
||||
graph = np.zeros(anim.shape[1], anim.shape[1])
|
||||
lengths = np.sum(anim.offsets**2.0, axis=1)**0.5 + 0.001
|
||||
|
||||
for i,p in enumerate(anim.parents):
|
||||
if p == -1: continue
|
||||
graph[i,p] = lengths[p]
|
||||
graph[p,i] = lengths[p]
|
||||
|
||||
return graph
|
||||
|
||||
|
||||
def distances(anim):
|
||||
"""
|
||||
Generates a distance matrix for
|
||||
pairwise joint distances along
|
||||
the skeletal structure
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
anim : Animation
|
||||
input animation
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
distances : (N, N) ndarray
|
||||
array of pairwise distances
|
||||
along skeletal structure
|
||||
from some joint N to some
|
||||
joint M
|
||||
"""
|
||||
|
||||
distances = np.zeros((anim.shape[1], anim.shape[1]))
|
||||
generated = distances.copy().astype(bool)
|
||||
|
||||
joint_lengths = np.sum(anim.offsets**2.0, axis=1)**0.5
|
||||
joint_children = children_list(anim)
|
||||
joint_parents = parents_list(anim)
|
||||
|
||||
def find_distance(distances, generated, prev, i, j):
|
||||
|
||||
""" If root, identity, or already generated, return """
|
||||
if j == -1: return (0.0, True)
|
||||
if j == i: return (0.0, True)
|
||||
if generated[i,j]: return (distances[i,j], True)
|
||||
|
||||
""" Find best distances along parents and children """
|
||||
par_dists = [(joint_lengths[j], find_distance(distances, generated, j, i, p)) for p in joint_parents[j] if p != prev]
|
||||
out_dists = [(joint_lengths[c], find_distance(distances, generated, j, i, c)) for c in joint_children[j] if c != prev]
|
||||
|
||||
""" Check valid distance and not dead end """
|
||||
par_dists = [a + d for (a, (d, f)) in par_dists if f]
|
||||
out_dists = [a + d for (a, (d, f)) in out_dists if f]
|
||||
|
||||
""" All dead ends """
|
||||
if (out_dists + par_dists) == []: return (0.0, False)
|
||||
|
||||
""" Get minimum path """
|
||||
dist = min(out_dists + par_dists)
|
||||
distances[i,j] = dist; distances[j,i] = dist
|
||||
generated[i,j] = True; generated[j,i] = True
|
||||
|
||||
for i in xrange(anim.shape[1]):
|
||||
for j in xrange(anim.shape[1]):
|
||||
find_distance(distances, generated, -1, i, j)
|
||||
|
||||
return distances
|
||||
|
||||
def edges(parents):
|
||||
"""
|
||||
Animation structure edges
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
parents : (J) ndarray
|
||||
parents array
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
edges : (M, 2) ndarray
|
||||
array of pairs where each
|
||||
pair contains two indices of a joints
|
||||
which corrisponds to an edge in the
|
||||
joint structure going from parent to child.
|
||||
"""
|
||||
|
||||
return np.array(list(zip(parents, joints(parents)))[1:])
|
||||
|
||||
|
||||
def incidence(parents):
|
||||
"""
|
||||
Incidence Matrix
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
parents : (J) ndarray
|
||||
parents array
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
incidence : (N, M) ndarray
|
||||
|
||||
Matrix of N joint positions by
|
||||
M edges which each entry is either
|
||||
1 or -1 and multiplication by the
|
||||
joint positions returns the an
|
||||
array of vectors along each edge
|
||||
of the structure
|
||||
"""
|
||||
|
||||
es = edges(parents)
|
||||
|
||||
inc = np.zeros((len(parents)-1, len(parents))).astype(np.int)
|
||||
for i, e in enumerate(es):
|
||||
inc[i,e[0]] = 1
|
||||
inc[i,e[1]] = -1
|
||||
|
||||
return inc.T
|
||||
|
||||
+471
@@ -0,0 +1,471 @@
|
||||
from Quaternions import Quaternions
|
||||
from Animation import Animation
|
||||
import re
|
||||
import numpy as np
|
||||
import sys
|
||||
|
||||
|
||||
channelmap = {
|
||||
'Xrotation': 'x',
|
||||
'Yrotation': 'y',
|
||||
'Zrotation': 'z'
|
||||
}
|
||||
|
||||
channelmap_inv = {
|
||||
'x': 'Xrotation',
|
||||
'y': 'Yrotation',
|
||||
'z': 'Zrotation',
|
||||
}
|
||||
|
||||
ordermap = {
|
||||
'x': 0,
|
||||
'y': 1,
|
||||
'z': 2,
|
||||
}
|
||||
|
||||
|
||||
def load(filename, start=None, end=None, order=None, world=False):
|
||||
"""
|
||||
Reads a BVH file and constructs an animation
|
||||
|
||||
Parameters
|
||||
----------
|
||||
filename: str
|
||||
File to be opened
|
||||
|
||||
start : int
|
||||
Optional Starting Frame
|
||||
|
||||
end : int
|
||||
Optional Ending Frame
|
||||
|
||||
order : str
|
||||
Optional Specifier for joint order.
|
||||
Given as string E.G 'xyz', 'zxy'
|
||||
|
||||
world : bool
|
||||
If set to true euler angles are applied
|
||||
together in world space rather than local
|
||||
space
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
(animation, joint_names, frametime)
|
||||
Tuple of loaded animation and joint names
|
||||
"""
|
||||
|
||||
f = open(filename, "r")
|
||||
|
||||
i = 0
|
||||
active = -1
|
||||
end_site = False
|
||||
|
||||
names = []
|
||||
orients = Quaternions.id(0)
|
||||
offsets = np.array([]).reshape((0, 3))
|
||||
parents = np.array([], dtype=int)
|
||||
|
||||
for line in f:
|
||||
|
||||
if "HIERARCHY" in line:
|
||||
continue
|
||||
if "MOTION" in line:
|
||||
continue
|
||||
|
||||
rmatch = re.match(r"ROOT (\w+)", line)
|
||||
if rmatch:
|
||||
names.append(rmatch.group(1))
|
||||
offsets = np.append(offsets, np.array([[0, 0, 0]]), axis=0)
|
||||
orients.qs = np.append(orients.qs, np.array([[1, 0, 0, 0]]), axis=0)
|
||||
parents = np.append(parents, active)
|
||||
active = (len(parents)-1)
|
||||
continue
|
||||
|
||||
if "{" in line:
|
||||
continue
|
||||
|
||||
if "}" in line:
|
||||
if end_site:
|
||||
end_site = False
|
||||
else:
|
||||
active = parents[active]
|
||||
continue
|
||||
|
||||
offmatch = re.match(r"\s*OFFSET\s+([\-\d\.e]+)\s+([\-\d\.e]+)\s+([\-\d\.e]+)", line)
|
||||
if offmatch:
|
||||
if not end_site:
|
||||
offsets[active] = np.array([list(map(float, offmatch.groups()))])
|
||||
continue
|
||||
|
||||
chanmatch = re.match(r"\s*CHANNELS\s+(\d+)", line)
|
||||
if chanmatch:
|
||||
channels = int(chanmatch.group(1))
|
||||
if order is None:
|
||||
channelis = 0 if channels == 3 else 3
|
||||
channelie = 3 if channels == 3 else 6
|
||||
parts = line.split()[2+channelis:2+channelie]
|
||||
if any([p not in channelmap for p in parts]):
|
||||
continue
|
||||
order = "".join([channelmap[p] for p in parts])
|
||||
continue
|
||||
|
||||
jmatch = re.match("\s*JOINT\s+(\w+)", line)
|
||||
if jmatch:
|
||||
names.append(jmatch.group(1))
|
||||
offsets = np.append(offsets, np.array([[0, 0, 0]]), axis=0)
|
||||
orients.qs = np.append(orients.qs, np.array([[1, 0, 0, 0]]), axis=0)
|
||||
parents = np.append(parents, active)
|
||||
active = (len(parents)-1)
|
||||
continue
|
||||
|
||||
if "End Site" in line:
|
||||
end_site = True
|
||||
continue
|
||||
|
||||
fmatch = re.match("\s*Frames:\s+(\d+)", line)
|
||||
if fmatch:
|
||||
if start and end:
|
||||
fnum = (end - start)-1
|
||||
else:
|
||||
fnum = int(fmatch.group(1))
|
||||
jnum = len(parents)
|
||||
# result: [fnum, J, 3]
|
||||
positions = offsets[np.newaxis].repeat(fnum, axis=0)
|
||||
# result: [fnum, len(orients), 3]
|
||||
rotations = np.zeros((fnum, len(orients), 3))
|
||||
continue
|
||||
|
||||
fmatch = re.match("\s*Frame Time:\s+([\d\.]+)", line)
|
||||
if fmatch:
|
||||
frametime = float(fmatch.group(1))
|
||||
continue
|
||||
|
||||
if (start and end) and (i < start or i >= end-1):
|
||||
i += 1
|
||||
continue
|
||||
|
||||
dmatch = line.strip().split()
|
||||
if dmatch:
|
||||
data_block = np.array(list(map(float, dmatch)))
|
||||
N = len(parents)
|
||||
fi = i - start if start else i
|
||||
if channels == 3:
|
||||
# This should be root positions[0:1] & all rotations
|
||||
positions[fi, 0:1] = data_block[0:3]
|
||||
rotations[fi, :] = data_block[3:].reshape(N, 3)
|
||||
elif channels == 6:
|
||||
data_block = data_block.reshape(N, 6)
|
||||
# fill in all positions
|
||||
positions[fi, :] = data_block[:, 0:3]
|
||||
rotations[fi, :] = data_block[:, 3:6]
|
||||
elif channels == 9:
|
||||
positions[fi, 0] = data_block[0:3]
|
||||
data_block = data_block[3:].reshape(N-1, 9)
|
||||
rotations[fi, 1:] = data_block[:, 3:6]
|
||||
positions[fi, 1:] += data_block[:, 0:3] * data_block[:, 6:9]
|
||||
else:
|
||||
raise Exception("Too many channels! %i" % channels)
|
||||
|
||||
i += 1
|
||||
|
||||
f.close()
|
||||
|
||||
rotations = Quaternions.from_euler(np.radians(rotations), order=order, world=world)
|
||||
|
||||
return (Animation(rotations, positions, orients, offsets, parents), names, frametime)
|
||||
|
||||
|
||||
def load_bfa(filename, start=None, end=None, order=None, world=False):
|
||||
"""
|
||||
Reads a BVH file and constructs an animation
|
||||
|
||||
!!! Read from bfa, will replace the end sites of arms by two joints (w/ unit rotation)
|
||||
|
||||
Parameters
|
||||
----------
|
||||
filename: str
|
||||
File to be opened
|
||||
|
||||
start : int
|
||||
Optional Starting Frame
|
||||
|
||||
end : int
|
||||
Optional Ending Frame
|
||||
|
||||
order : str
|
||||
Optional Specifier for joint order.
|
||||
Given as string E.G 'xyz', 'zxy'
|
||||
|
||||
world : bool
|
||||
If set to true euler angles are applied
|
||||
together in world space rather than local
|
||||
space
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
(animation, joint_names, frametime)
|
||||
Tuple of loaded animation and joint names
|
||||
"""
|
||||
|
||||
f = open(filename, "r")
|
||||
|
||||
i = 0
|
||||
active = -1
|
||||
end_site = False
|
||||
|
||||
hand_idx = [9, 14]
|
||||
|
||||
names = []
|
||||
orients = Quaternions.id(0)
|
||||
offsets = np.array([]).reshape((0, 3))
|
||||
parents = np.array([], dtype=int)
|
||||
|
||||
for line in f:
|
||||
|
||||
if "HIERARCHY" in line:
|
||||
continue
|
||||
if "MOTION" in line:
|
||||
continue
|
||||
|
||||
rmatch = re.match(r"ROOT (\w+)", line)
|
||||
if rmatch:
|
||||
names.append(rmatch.group(1))
|
||||
offsets = np.append(offsets, np.array([[0, 0, 0]]), axis=0)
|
||||
orients.qs = np.append(orients.qs, np.array([[1, 0, 0, 0]]), axis=0)
|
||||
parents = np.append(parents, active)
|
||||
active = (len(parents)-1)
|
||||
continue
|
||||
|
||||
if "{" in line:
|
||||
continue
|
||||
|
||||
if "}" in line:
|
||||
if end_site:
|
||||
end_site = False
|
||||
else:
|
||||
active = parents[active]
|
||||
continue
|
||||
|
||||
offmatch = re.match(r"\s*OFFSET\s+([\-\d\.e]+)\s+([\-\d\.e]+)\s+([\-\d\.e]+)", line)
|
||||
if offmatch:
|
||||
if not end_site:
|
||||
offsets[active] = np.array([list(map(float, offmatch.groups()))])
|
||||
"""
|
||||
else:
|
||||
print("active = ", active)
|
||||
if active in hand_idx:
|
||||
offsets[active] = np.array([list(map(float, offmatch.groups()))])
|
||||
"""
|
||||
continue
|
||||
|
||||
chanmatch = re.match(r"\s*CHANNELS\s+(\d+)", line)
|
||||
if chanmatch:
|
||||
channels = int(chanmatch.group(1))
|
||||
if order is None:
|
||||
channelis = 0 if channels == 3 else 3
|
||||
channelie = 3 if channels == 3 else 6
|
||||
parts = line.split()[2+channelis:2+channelie]
|
||||
if any([p not in channelmap for p in parts]):
|
||||
continue
|
||||
order = "".join([channelmap[p] for p in parts])
|
||||
continue
|
||||
|
||||
jmatch = re.match("\s*JOINT\s+(\w+)", line)
|
||||
if jmatch:
|
||||
names.append(jmatch.group(1))
|
||||
offsets = np.append(offsets, np.array([[0, 0, 0]]), axis=0)
|
||||
orients.qs = np.append(orients.qs, np.array([[1, 0, 0, 0]]), axis=0)
|
||||
parents = np.append(parents, active)
|
||||
active = (len(parents)-1)
|
||||
continue
|
||||
|
||||
if "End Site" in line:
|
||||
if active + 1 in hand_idx:
|
||||
print("parent:", names[-1])
|
||||
name = "LeftHandIndex" if active + 1 == hand_idx[0] else "RightHandIndex"
|
||||
names.append(name)
|
||||
offsets = np.append(offsets, np.array([[0, 0, 0]]), axis=0)
|
||||
orients.qs = np.append(orients.qs, np.array([[1, 0, 0, 0]]), axis=0)
|
||||
parents = np.append(parents, active)
|
||||
active = (len(parents)-1)
|
||||
else:
|
||||
end_site = True
|
||||
continue
|
||||
|
||||
fmatch = re.match("\s*Frames:\s+(\d+)", line)
|
||||
if fmatch:
|
||||
if start and end:
|
||||
fnum = (end - start)-1
|
||||
else:
|
||||
fnum = int(fmatch.group(1))
|
||||
jnum = len(parents)
|
||||
# result: [fnum, J, 3]
|
||||
positions = offsets[np.newaxis].repeat(fnum, axis=0)
|
||||
# result: [fnum, len(orients), 3]
|
||||
rotations = np.zeros((fnum, len(orients), 3))
|
||||
continue
|
||||
|
||||
fmatch = re.match("\s*Frame Time:\s+([\d\.]+)", line)
|
||||
if fmatch:
|
||||
frametime = float(fmatch.group(1))
|
||||
continue
|
||||
|
||||
if (start and end) and (i < start or i >= end-1):
|
||||
i += 1
|
||||
continue
|
||||
|
||||
dmatch = line.strip().split()
|
||||
if dmatch:
|
||||
data_block = np.array(list(map(float, dmatch)))
|
||||
N = len(parents)
|
||||
fi = i - start if start else i
|
||||
if channels == 3:
|
||||
# This should be root positions[0:1] & all rotations
|
||||
positions[fi, 0:1] = data_block[0:3]
|
||||
tmp = data_block[3:].reshape(N - 2, 3)
|
||||
tmp = np.concatenate([tmp[:hand_idx[0]],
|
||||
np.array([[0, 0, 0]]),
|
||||
tmp[hand_idx[0]: hand_idx[1] - 1],
|
||||
np.array([[0, 0, 0]]),
|
||||
tmp[hand_idx[1] - 1:]], axis=0)
|
||||
rotations[fi, :] = tmp.reshape(N, 3)
|
||||
elif channels == 6:
|
||||
data_block = data_block.reshape(N, 6)
|
||||
# fill in all positions
|
||||
positions[fi, :] = data_block[:, 0:3]
|
||||
rotations[fi, :] = data_block[:, 3:6]
|
||||
elif channels == 9:
|
||||
positions[fi, 0] = data_block[0:3]
|
||||
data_block = data_block[3:].reshape(N-1, 9)
|
||||
rotations[fi, 1:] = data_block[:, 3:6]
|
||||
positions[fi, 1:] += data_block[:, 0:3] * data_block[:, 6:9]
|
||||
else:
|
||||
raise Exception("Too many channels! %i" % channels)
|
||||
|
||||
i += 1
|
||||
|
||||
f.close()
|
||||
|
||||
rotations = Quaternions.from_euler(np.radians(rotations), order=order, world=world)
|
||||
|
||||
return (Animation(rotations, positions, orients, offsets, parents), names, frametime)
|
||||
|
||||
|
||||
def save(filename, anim, names=None, frametime=1.0/24.0, order='zyx', positions=False, orients=True):
|
||||
"""
|
||||
Saves an Animation to file as BVH
|
||||
|
||||
Parameters
|
||||
----------
|
||||
filename: str
|
||||
File to be saved to
|
||||
|
||||
anim : Animation
|
||||
Animation to save
|
||||
|
||||
names : [str]
|
||||
List of joint names
|
||||
|
||||
order : str
|
||||
Optional Specifier for joint order.
|
||||
Given as string E.G 'xyz', 'zxy'
|
||||
|
||||
frametime : float
|
||||
Optional Animation Frame time
|
||||
|
||||
positions : bool
|
||||
Optional specfier to save bone
|
||||
positions for each frame
|
||||
|
||||
orients : bool
|
||||
Multiply joint orients to the rotations
|
||||
before saving.
|
||||
|
||||
"""
|
||||
|
||||
if names is None:
|
||||
names = ["joint_" + str(i) for i in range(len(anim.parents))]
|
||||
|
||||
with open(filename, 'w') as f:
|
||||
|
||||
t = ""
|
||||
f.write("%sHIERARCHY\n" % t)
|
||||
f.write("%sROOT %s\n" % (t, names[0]))
|
||||
f.write("%s{\n" % t)
|
||||
t += '\t'
|
||||
|
||||
f.write("%sOFFSET %f %f %f\n" % (t, anim.offsets[0, 0], anim.offsets[0, 1], anim.offsets[0, 2]))
|
||||
f.write("%sCHANNELS 6 Xposition Yposition Zposition %s %s %s \n" %
|
||||
(t, channelmap_inv[order[0]], channelmap_inv[order[1]], channelmap_inv[order[2]]))
|
||||
|
||||
for i in range(anim.shape[1]):
|
||||
if anim.parents[i] == 0:
|
||||
t = save_joint(f, anim, names, t, i, order=order, positions=positions)
|
||||
|
||||
t = t[:-1]
|
||||
f.write("%s}\n" % t)
|
||||
|
||||
f.write("MOTION\n")
|
||||
f.write("Frames: %i\n" % anim.shape[0])
|
||||
f.write("Frame Time: %f\n" % frametime)
|
||||
|
||||
# if orients:
|
||||
# rots = np.degrees((-anim.orients[np.newaxis] * anim.rotations).euler(order=order[::-1]))
|
||||
# else:
|
||||
# rots = np.degrees(anim.rotations.euler(order=order[::-1]))
|
||||
rots = np.degrees(anim.rotations.euler(order=order[::-1]))
|
||||
poss = anim.positions
|
||||
|
||||
for i in range(anim.shape[0]):
|
||||
for j in range(anim.shape[1]):
|
||||
|
||||
if positions or j == 0:
|
||||
|
||||
f.write("%f %f %f %f %f %f " % (
|
||||
poss[i, j, 0], poss[i, j, 1], poss[i, j, 2],
|
||||
rots[i, j, ordermap[order[0]]], rots[i, j, ordermap[order[1]]], rots[i, j, ordermap[order[2]]]))
|
||||
|
||||
else:
|
||||
|
||||
f.write("%f %f %f " % (
|
||||
rots[i, j, ordermap[order[0]]], rots[i, j, ordermap[order[1]]], rots[i, j, ordermap[order[2]]]))
|
||||
|
||||
f.write("\n")
|
||||
|
||||
|
||||
def save_joint(f, anim, names, t, i, order='zyx', positions=False):
|
||||
|
||||
f.write("%sJOINT %s\n" % (t, names[i]))
|
||||
f.write("%s{\n" % t)
|
||||
t += '\t'
|
||||
|
||||
f.write("%sOFFSET %f %f %f\n" % (t, anim.offsets[i, 0], anim.offsets[i, 1], anim.offsets[i, 2]))
|
||||
|
||||
if positions:
|
||||
f.write("%sCHANNELS 6 Xposition Yposition Zposition %s %s %s \n" % (t,
|
||||
channelmap_inv[order[0]], channelmap_inv[order[1]], channelmap_inv[order[2]]))
|
||||
else:
|
||||
f.write("%sCHANNELS 3 %s %s %s\n" % (t,
|
||||
channelmap_inv[order[0]], channelmap_inv[order[1]], channelmap_inv[order[2]]))
|
||||
|
||||
end_site = True
|
||||
|
||||
for j in range(anim.shape[1]):
|
||||
if anim.parents[j] == i:
|
||||
t = save_joint(f, anim, names, t, j, order=order, positions=positions)
|
||||
end_site = False
|
||||
|
||||
if end_site:
|
||||
f.write("%sEnd Site\n" % t)
|
||||
f.write("%s{\n" % t)
|
||||
t += '\t'
|
||||
f.write("%sOFFSET %f %f %f\n" % (t, 0.0, 0.0, 0.0))
|
||||
t = t[:-1]
|
||||
f.write("%s}\n" % t)
|
||||
|
||||
t = t[:-1]
|
||||
f.write("%s}\n" % t)
|
||||
|
||||
return t
|
||||
@@ -0,0 +1,496 @@
|
||||
import numpy as np
|
||||
import scipy.linalg as linalg
|
||||
|
||||
import Animation
|
||||
import AnimationStructure
|
||||
|
||||
from Quaternions import Quaternions
|
||||
|
||||
|
||||
class BasicInverseKinematics:
|
||||
"""
|
||||
Basic Inverse Kinematics Solver
|
||||
|
||||
This is an extremely simple full body IK
|
||||
solver.
|
||||
|
||||
It works given the following conditions:
|
||||
|
||||
* All joint targets must be specified
|
||||
* All joint targets must be in reach
|
||||
* All joint targets must not differ
|
||||
extremely from the starting pose
|
||||
* No bone length constraints can be violated
|
||||
* The root translation and rotation are
|
||||
set to good initial values
|
||||
|
||||
It works under the observation that if the
|
||||
_directions_ the joints are pointing toward
|
||||
match the _directions_ of the vectors between
|
||||
the target joints then the pose should match
|
||||
that of the target pose.
|
||||
|
||||
Therefore it iterates over joints rotating
|
||||
each joint such that the vectors between it
|
||||
and it's children match that of the target
|
||||
positions.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
animation : Animation
|
||||
animation input
|
||||
|
||||
positions : (F, J, 3) ndarray
|
||||
target positions for each frame F
|
||||
and each joint J
|
||||
|
||||
iterations : int
|
||||
Optional number of iterations.
|
||||
If the above conditions are met
|
||||
1 iteration should be enough,
|
||||
therefore the default is 1
|
||||
|
||||
silent : bool
|
||||
Optional if to suppress output
|
||||
defaults to False
|
||||
"""
|
||||
|
||||
def __init__(self, animation, positions, iterations=1, silent=True):
|
||||
|
||||
self.animation = animation
|
||||
self.positions = positions
|
||||
self.iterations = iterations
|
||||
self.silent = silent
|
||||
|
||||
def __call__(self):
|
||||
|
||||
children = AnimationStructure.children_list(self.animation.parents)
|
||||
|
||||
for i in range(self.iterations):
|
||||
|
||||
for j in AnimationStructure.joints(self.animation.parents):
|
||||
|
||||
c = np.array(children[j])
|
||||
if len(c) == 0:
|
||||
continue
|
||||
|
||||
anim_transforms = Animation.transforms_global(self.animation)
|
||||
anim_positions = anim_transforms[:, :, :3, 3]
|
||||
anim_rotations = Quaternions.from_transforms(anim_transforms)
|
||||
|
||||
jdirs = anim_positions[:, c] - anim_positions[:, np.newaxis, j]
|
||||
ddirs = self.positions[:, c] - anim_positions[:, np.newaxis, j]
|
||||
|
||||
jsums = np.sqrt(np.sum(jdirs**2.0, axis=-1)) + 1e-10
|
||||
dsums = np.sqrt(np.sum(ddirs**2.0, axis=-1)) + 1e-10
|
||||
|
||||
jdirs = jdirs / jsums[:, :, np.newaxis]
|
||||
ddirs = ddirs / dsums[:, :, np.newaxis]
|
||||
|
||||
angles = np.arccos(np.sum(jdirs * ddirs, axis=2).clip(-1, 1))
|
||||
axises = np.cross(jdirs, ddirs)
|
||||
axises = -anim_rotations[:, j, np.newaxis] * axises
|
||||
|
||||
rotations = Quaternions.from_angle_axis(angles, axises)
|
||||
|
||||
if rotations.shape[1] == 1:
|
||||
averages = rotations[:, 0]
|
||||
else:
|
||||
averages = Quaternions.exp(rotations.log().mean(axis=-2))
|
||||
|
||||
self.animation.rotations[:, j] = self.animation.rotations[:, j] * averages
|
||||
|
||||
if not self.silent:
|
||||
anim_positions = Animation.positions_global(self.animation)
|
||||
error = np.mean(np.sum((anim_positions - self.positions)**2.0, axis=-1)**0.5, axis=-1)
|
||||
print('[BasicInverseKinematics] Iteration %i Error: %f' % (i+1, error))
|
||||
|
||||
return self.animation
|
||||
|
||||
|
||||
class JacobianInverseKinematics:
|
||||
"""
|
||||
Jacobian Based Full Body IK Solver
|
||||
|
||||
This is a full body IK solver which
|
||||
uses the dampened least squares inverse
|
||||
jacobian method.
|
||||
|
||||
It should remain fairly stable and effective
|
||||
even for joint positions which are out of
|
||||
reach and it can also take any number of targets
|
||||
to treat as end effectors.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
animation : Animation
|
||||
animation to solve inverse problem on
|
||||
|
||||
targets : {int : (F, 3) ndarray}
|
||||
Dictionary of target positions for each
|
||||
frame F, mapping joint index to
|
||||
a target position
|
||||
|
||||
references : (F, 3)
|
||||
Optional list of J joint position
|
||||
references for which the result
|
||||
should bias toward
|
||||
|
||||
iterations : int
|
||||
Optional number of iterations to
|
||||
compute. More iterations results in
|
||||
better accuracy but takes longer to
|
||||
compute. Default is 10.
|
||||
|
||||
recalculate : bool
|
||||
Optional if to recalcuate jacobian
|
||||
each iteration. Gives better accuracy
|
||||
but slower to compute. Defaults to True
|
||||
|
||||
damping : float
|
||||
Optional damping constant. Higher
|
||||
damping increases stability but
|
||||
requires more iterations to converge.
|
||||
Defaults to 5.0
|
||||
|
||||
secondary : float
|
||||
Force, or bias toward secondary target.
|
||||
Defaults to 0.25
|
||||
|
||||
silent : bool
|
||||
Optional if to suppress output
|
||||
defaults to False
|
||||
"""
|
||||
|
||||
def __init__(self, animation, targets,
|
||||
references=None, iterations=10,
|
||||
recalculate=True, damping=2.0,
|
||||
secondary=0.25, translate=False,
|
||||
silent=False, weights=None,
|
||||
weights_translate=None):
|
||||
|
||||
self.animation = animation
|
||||
self.targets = targets
|
||||
self.references = references
|
||||
|
||||
self.iterations = iterations
|
||||
self.recalculate = recalculate
|
||||
self.damping = damping
|
||||
self.secondary = secondary
|
||||
self.translate = translate
|
||||
self.silent = silent
|
||||
self.weights = weights
|
||||
self.weights_translate = weights_translate
|
||||
|
||||
def cross(self, a, b):
|
||||
o = np.empty(b.shape)
|
||||
o[..., 0] = a[..., 1]*b[..., 2] - a[..., 2]*b[..., 1]
|
||||
o[..., 1] = a[..., 2]*b[..., 0] - a[..., 0]*b[..., 2]
|
||||
o[..., 2] = a[..., 0]*b[..., 1] - a[..., 1]*b[..., 0]
|
||||
return o
|
||||
|
||||
def jacobian(self, x, fp, fr, ts, dsc, tdsc):
|
||||
""" Find parent rotations """
|
||||
prs = fr[:, self.animation.parents]
|
||||
prs[:, 0] = Quaternions.id((1))
|
||||
|
||||
""" Find global positions of target joints """
|
||||
tps = fp[:, np.array(list(ts.keys()))]
|
||||
|
||||
""" Get partial rotations """
|
||||
qys = Quaternions.from_angle_axis(x[:, 1:prs.shape[1]*3:3], np.array([[[0, 1, 0]]]))
|
||||
qzs = Quaternions.from_angle_axis(x[:, 2:prs.shape[1]*3:3], np.array([[[0, 0, 1]]]))
|
||||
|
||||
""" Find axis of rotations """
|
||||
es = np.empty((len(x), fr.shape[1]*3, 3))
|
||||
es[:, 0::3] = ((prs * qzs) * qys) * np.array([[[1, 0, 0]]])
|
||||
es[:, 1::3] = ((prs * qzs) * np.array([[[0, 1, 0]]]))
|
||||
es[:, 2::3] = ((prs * np.array([[[0, 0, 1]]])))
|
||||
|
||||
""" Construct Jacobian """
|
||||
j = fp.repeat(3, axis=1)
|
||||
j = dsc[np.newaxis, :, :, np.newaxis] * (tps[:, np.newaxis, :] - j[:, :, np.newaxis])
|
||||
j = self.cross(es[:, :, np.newaxis, :], j)
|
||||
j = np.swapaxes(j.reshape((len(x), fr.shape[1]*3, len(ts)*3)), 1, 2)
|
||||
|
||||
if self.translate:
|
||||
|
||||
es = np.empty((len(x), fr.shape[1]*3, 3))
|
||||
es[:, 0::3] = prs * np.array([[[1, 0, 0]]])
|
||||
es[:, 1::3] = prs * np.array([[[0, 1, 0]]])
|
||||
es[:, 2::3] = prs * np.array([[[0, 0, 1]]])
|
||||
|
||||
jt = tdsc[np.newaxis, :, :, np.newaxis] * es[:, :, np.newaxis, :].repeat(tps.shape[1], axis=2)
|
||||
jt = np.swapaxes(jt.reshape((len(x), fr.shape[1]*3, len(ts)*3)), 1, 2)
|
||||
|
||||
j = np.concatenate([j, jt], axis=-1)
|
||||
|
||||
return j
|
||||
|
||||
# @profile(immediate=True)
|
||||
def __call__(self, descendants=None, gamma=1.0):
|
||||
|
||||
self.descendants = descendants
|
||||
|
||||
""" Calculate Masses """
|
||||
if self.weights is None:
|
||||
self.weights = np.ones(self.animation.shape[1])
|
||||
|
||||
if self.weights_translate is None:
|
||||
self.weights_translate = np.ones(self.animation.shape[1])
|
||||
|
||||
""" Calculate Descendants """
|
||||
if self.descendants is None:
|
||||
self.descendants = AnimationStructure.descendants_mask(self.animation.parents)
|
||||
|
||||
self.tdescendants = np.eye(self.animation.shape[1]) + self.descendants
|
||||
|
||||
self.first_descendants = self.descendants[:, np.array(list(self.targets.keys()))].repeat(3, axis=0).astype(int)
|
||||
self.first_tdescendants = self.tdescendants[:, np.array(list(self.targets.keys()))].repeat(3, axis=0).astype(int)
|
||||
|
||||
""" Calculate End Effectors """
|
||||
self.endeff = np.array(list(self.targets.values()))
|
||||
self.endeff = np.swapaxes(self.endeff, 0, 1)
|
||||
|
||||
if not self.references is None:
|
||||
self.second_descendants = self.descendants.repeat(3, axis=0).astype(int)
|
||||
self.second_tdescendants = self.tdescendants.repeat(3, axis=0).astype(int)
|
||||
self.second_targets = dict([(i, self.references[:, i]) for i in range(self.references.shape[1])])
|
||||
|
||||
nf = len(self.animation)
|
||||
nj = self.animation.shape[1]
|
||||
|
||||
if not self.silent:
|
||||
gp = Animation.positions_global(self.animation)
|
||||
gp = gp[:, np.array(list(self.targets.keys()))]
|
||||
error = np.mean(np.sqrt(np.sum((self.endeff - gp)**2.0, axis=2)))
|
||||
print('[JacobianInverseKinematics] Start | Error: %f' % error)
|
||||
|
||||
for i in range(self.iterations):
|
||||
|
||||
""" Get Global Rotations & Positions """
|
||||
gt = Animation.transforms_global(self.animation)
|
||||
gp = gt[:, :, :, 3]
|
||||
gp = gp[:, :, :3] / gp[:, :, 3, np.newaxis]
|
||||
gr = Quaternions.from_transforms(gt)
|
||||
|
||||
x = self.animation.rotations.euler().reshape(nf, -1)
|
||||
w = self.weights.repeat(3)
|
||||
|
||||
if self.translate:
|
||||
x = np.hstack([x, self.animation.positions.reshape(nf, -1)])
|
||||
w = np.hstack([w, self.weights_translate.repeat(3)])
|
||||
|
||||
""" Generate Jacobian """
|
||||
if self.recalculate or i == 0:
|
||||
j = self.jacobian(x, gp, gr, self.targets, self.first_descendants, self.first_tdescendants)
|
||||
|
||||
""" Update Variables """
|
||||
l = self.damping * (1.0 / (w + 0.001))
|
||||
d = (l*l) * np.eye(x.shape[1])
|
||||
e = gamma * (self.endeff.reshape(nf, -1) - gp[:, np.array(list(self.targets.keys()))].reshape(nf, -1))
|
||||
|
||||
j = np.nan_to_num(j)
|
||||
e = np.nan_to_num(e)
|
||||
x += np.array(list(map(lambda jf, ef:
|
||||
linalg.lu_solve(linalg.lu_factor(jf.T.dot(jf) + d), jf.T.dot(ef)), j, e)))
|
||||
|
||||
""" Generate Secondary Jacobian """
|
||||
if self.references is not None:
|
||||
|
||||
ns = np.array(list(map(lambda jf:
|
||||
np.eye(x.shape[1]) - linalg.solve(jf.T.dot(jf) + d, jf.T.dot(jf)), j)))
|
||||
|
||||
if self.recalculate or i == 0:
|
||||
j2 = self.jacobian(x, gp, gr, self.second_targets, self.second_descendants, self.second_tdescendants)
|
||||
|
||||
e2 = self.secondary * (self.references.reshape(nf, -1) - gp.reshape(nf, -1))
|
||||
|
||||
x += np.array(list(map(lambda nsf, j2f, e2f:
|
||||
nsf.dot(linalg.lu_solve(linalg.lu_factor(j2f.T.dot(j2f) + d), j2f.T.dot(e2f))), ns, j2, e2)))
|
||||
|
||||
""" Set Back Rotations / Translations """
|
||||
self.animation.rotations = Quaternions.from_euler(
|
||||
x[:, :nj*3].reshape((nf, nj, 3)), order='xyz', world=True)
|
||||
|
||||
if self.translate:
|
||||
self.animation.positions = x[:, nj*3:].reshape((nf, nj, 3))
|
||||
|
||||
""" Generate Error """
|
||||
|
||||
if not self.silent:
|
||||
gp = Animation.positions_global(self.animation)
|
||||
gp = gp[:, np.array(list(self.targets.keys()))]
|
||||
error = np.mean(np.sum((self.endeff - gp)**2.0, axis=2)**0.5)
|
||||
print('[JacobianInverseKinematics] Iteration %i | Error: %f' % (i+1, error))
|
||||
|
||||
|
||||
class BasicJacobianIK:
|
||||
"""
|
||||
Same interface as BasicInverseKinematics
|
||||
but uses the Jacobian IK Solver Instead
|
||||
"""
|
||||
|
||||
def __init__(self, animation, positions, iterations=10, silent=True, **kw):
|
||||
|
||||
targets = dict([(i, positions[:, i]) for i in range(positions.shape[1])])
|
||||
self.ik = JacobianInverseKinematics(animation, targets, iterations=iterations, silent=silent, **kw)
|
||||
|
||||
def __call__(self, **kw):
|
||||
return self.ik(**kw)
|
||||
|
||||
|
||||
class ICP:
|
||||
|
||||
def __init__(self,
|
||||
anim, rest, weights, mesh, goal,
|
||||
find_closest=True, damping=10,
|
||||
iterations=10, silent=True,
|
||||
translate=True, recalculate=True,
|
||||
weights_translate=None):
|
||||
|
||||
self.animation = anim
|
||||
self.rest = rest
|
||||
self.vweights = weights
|
||||
self.mesh = mesh
|
||||
self.goal = goal
|
||||
self.find_closest = find_closest
|
||||
self.iterations = iterations
|
||||
self.silent = silent
|
||||
self.translate = translate
|
||||
self.damping = damping
|
||||
self.weights = None
|
||||
self.weights_translate = weights_translate
|
||||
self.recalculate = recalculate
|
||||
|
||||
def cross(self, a, b):
|
||||
o = np.empty(b.shape)
|
||||
o[..., 0] = a[..., 1]*b[..., 2] - a[..., 2]*b[..., 1]
|
||||
o[..., 1] = a[..., 2]*b[..., 0] - a[..., 0]*b[..., 2]
|
||||
o[..., 2] = a[..., 0]*b[..., 1] - a[..., 1]*b[..., 0]
|
||||
return o
|
||||
|
||||
def jacobian(self, x, fp, fr, goal, weights, des_r, des_t):
|
||||
""" Find parent rotations """
|
||||
prs = fr[:, self.animation.parents]
|
||||
prs[:, 0] = Quaternions.id((1))
|
||||
|
||||
""" Get partial rotations """
|
||||
qys = Quaternions.from_angle_axis(x[:, 1:prs.shape[1]*3:3], np.array([[[0, 1, 0]]]))
|
||||
qzs = Quaternions.from_angle_axis(x[:, 2:prs.shape[1]*3:3], np.array([[[0, 0, 1]]]))
|
||||
|
||||
""" Find axis of rotations """
|
||||
es = np.empty((len(x), fr.shape[1]*3, 3))
|
||||
es[:, 0::3] = ((prs * qzs) * qys) * np.array([[[1, 0, 0]]])
|
||||
es[:, 1::3] = ((prs * qzs) * np.array([[[0, 1, 0]]]))
|
||||
es[:, 2::3] = ((prs * np.array([[[0, 0, 1]]])))
|
||||
|
||||
""" Construct Jacobian """
|
||||
j = fp.repeat(3, axis=1)
|
||||
j = des_r[np.newaxis, :, :, :, np.newaxis] * (goal[:, np.newaxis, :, np.newaxis] - j[:, :, np.newaxis, np.newaxis])
|
||||
j = np.sum(j * weights[np.newaxis, np.newaxis, :, :, np.newaxis], 3)
|
||||
j = self.cross(es[:, :, np.newaxis, :], j)
|
||||
j = np.swapaxes(j.reshape((len(x), fr.shape[1]*3, goal.shape[1]*3)), 1, 2)
|
||||
|
||||
if self.translate:
|
||||
|
||||
es = np.empty((len(x), fr.shape[1]*3, 3))
|
||||
es[:, 0::3] = prs * np.array([[[1, 0, 0]]])
|
||||
es[:, 1::3] = prs * np.array([[[0, 1, 0]]])
|
||||
es[:, 2::3] = prs * np.array([[[0, 0, 1]]])
|
||||
|
||||
jt = des_t[np.newaxis, :, :, :, np.newaxis] * es[:, :, np.newaxis, np.newaxis, :].repeat(goal.shape[1], axis=2)
|
||||
jt = np.sum(jt * weights[np.newaxis, np.newaxis, :, :, np.newaxis], 3)
|
||||
jt = np.swapaxes(jt.reshape((len(x), fr.shape[1]*3, goal.shape[1]*3)), 1, 2)
|
||||
|
||||
j = np.concatenate([j, jt], axis=-1)
|
||||
|
||||
return j
|
||||
|
||||
# @profile(immediate=True)
|
||||
def __call__(self, descendants=None, maxjoints=4, gamma=1.0, transpose=False):
|
||||
""" Calculate Masses """
|
||||
if self.weights is None:
|
||||
self.weights = np.ones(self.animation.shape[1])
|
||||
|
||||
if self.weights_translate is None:
|
||||
self.weights_translate = np.ones(self.animation.shape[1])
|
||||
|
||||
nf = len(self.animation)
|
||||
nj = self.animation.shape[1]
|
||||
nv = self.goal.shape[1]
|
||||
|
||||
weightids = np.argsort(-self.vweights, axis=1)[:, :maxjoints]
|
||||
weightvls = np.array(list(map(lambda w, i: w[i], self.vweights, weightids)))
|
||||
weightvls = weightvls / weightvls.sum(axis=1)[..., np.newaxis]
|
||||
|
||||
if descendants is None:
|
||||
self.descendants = AnimationStructure.descendants_mask(self.animation.parents)
|
||||
else:
|
||||
self.descendants = descendants
|
||||
|
||||
des_r = np.eye(nj) + self.descendants
|
||||
des_r = des_r[:, weightids].repeat(3, axis=0)
|
||||
|
||||
des_t = np.eye(nj) + self.descendants
|
||||
des_t = des_t[:, weightids].repeat(3, axis=0)
|
||||
|
||||
if not self.silent:
|
||||
curr = Animation.skin(self.animation, self.rest, self.vweights, self.mesh, maxjoints=maxjoints)
|
||||
error = np.mean(np.sqrt(np.sum((curr - self.goal)**2.0, axis=-1)))
|
||||
print('[ICP] Start | Error: %f' % error)
|
||||
|
||||
for i in range(self.iterations):
|
||||
|
||||
""" Get Global Rotations & Positions """
|
||||
gt = Animation.transforms_global(self.animation)
|
||||
gp = gt[:, :, :, 3]
|
||||
gp = gp[:, :, :3] / gp[:, :, 3, np.newaxis]
|
||||
gr = Quaternions.from_transforms(gt)
|
||||
|
||||
x = self.animation.rotations.euler().reshape(nf, -1)
|
||||
w = self.weights.repeat(3)
|
||||
|
||||
if self.translate:
|
||||
x = np.hstack([x, self.animation.positions.reshape(nf, -1)])
|
||||
w = np.hstack([w, self.weights_translate.repeat(3)])
|
||||
|
||||
""" Get Current State """
|
||||
curr = Animation.skin(self.animation, self.rest, self.vweights, self.mesh, maxjoints=maxjoints)
|
||||
|
||||
""" Find Cloest Points """
|
||||
if self.find_closest:
|
||||
mapping = np.argmin(
|
||||
(curr[:, :, np.newaxis] -
|
||||
self.goal[:, np.newaxis, :])**2.0, axis=2)
|
||||
e = gamma * (np.array(list(map(lambda g, m: g[m], self.goal, mapping))) - curr).reshape(nf, -1)
|
||||
else:
|
||||
e = gamma * (self.goal - curr).reshape(nf, -1)
|
||||
|
||||
""" Generate Jacobian """
|
||||
if self.recalculate or i == 0:
|
||||
j = self.jacobian(x, gp, gr, self.goal, weightvls, des_r, des_t)
|
||||
|
||||
""" Update Variables """
|
||||
l = self.damping * (1.0 / (w + 1e-10))
|
||||
d = (l*l) * np.eye(x.shape[1])
|
||||
|
||||
if transpose:
|
||||
x += np.array(list(map(lambda jf, ef: jf.T.dot(ef), j, e)))
|
||||
else:
|
||||
x += np.array(list(map(lambda jf, ef:
|
||||
linalg.lu_solve(linalg.lu_factor(jf.T.dot(jf) + d), jf.T.dot(ef)), j, e)))
|
||||
|
||||
""" Set Back Rotations / Translations """
|
||||
self.animation.rotations = Quaternions.from_euler(
|
||||
x[:, :nj*3].reshape((nf, nj, 3)), order='xyz', world=True)
|
||||
|
||||
if self.translate:
|
||||
self.animation.positions = x[:, nj*3:].reshape((nf, nj, 3))
|
||||
|
||||
if not self.silent:
|
||||
curr = Animation.skin(self.animation, self.rest, self.vweights, self.mesh)
|
||||
error = np.mean(np.sqrt(np.sum((curr - self.goal)**2.0, axis=-1)))
|
||||
print('[ICP] Iteration %i | Error: %f' % (i+1, error))
|
||||
@@ -0,0 +1,91 @@
|
||||
import numpy as np
|
||||
|
||||
from Quaternions import Quaternions
|
||||
|
||||
|
||||
class Pivots:
|
||||
"""
|
||||
Pivots is an ndarray of angular rotations
|
||||
|
||||
This wrapper provides some functions for
|
||||
working with pivots.
|
||||
|
||||
These are particularly useful as a number
|
||||
of atomic operations (such as adding or
|
||||
subtracting) cannot be achieved using
|
||||
the standard arithmatic and need to be
|
||||
defined differently to work correctly
|
||||
"""
|
||||
|
||||
def __init__(self, ps): self.ps = np.array(ps)
|
||||
def __str__(self): return "Pivots(" + str(self.ps) + ")"
|
||||
def __repr__(self): return "Pivots(" + repr(self.ps) + ")"
|
||||
|
||||
def __add__(self, other): return Pivots(np.arctan2(np.sin(self.ps + other.ps), np.cos(self.ps + other.ps)))
|
||||
def __sub__(self, other): return Pivots(np.arctan2(np.sin(self.ps - other.ps), np.cos(self.ps - other.ps)))
|
||||
def __mul__(self, other): return Pivots(self.ps * other.ps)
|
||||
def __div__(self, other): return Pivots(self.ps / other.ps)
|
||||
def __mod__(self, other): return Pivots(self.ps % other.ps)
|
||||
def __pow__(self, other): return Pivots(self.ps ** other.ps)
|
||||
|
||||
def __lt__(self, other): return self.ps < other.ps
|
||||
def __le__(self, other): return self.ps <= other.ps
|
||||
def __eq__(self, other): return self.ps == other.ps
|
||||
def __ne__(self, other): return self.ps != other.ps
|
||||
def __ge__(self, other): return self.ps >= other.ps
|
||||
def __gt__(self, other): return self.ps > other.ps
|
||||
|
||||
def __abs__(self): return Pivots(abs(self.ps))
|
||||
def __neg__(self): return Pivots(-self.ps)
|
||||
|
||||
def __iter__(self): return iter(self.ps)
|
||||
def __len__(self): return len(self.ps)
|
||||
|
||||
def __getitem__(self, k): return Pivots(self.ps[k])
|
||||
def __setitem__(self, k, v): self.ps[k] = v.ps
|
||||
|
||||
def _ellipsis(self): return tuple(map(lambda x: slice(None), self.shape))
|
||||
|
||||
def quaternions(self, plane='xz'):
|
||||
fa = self._ellipsis()
|
||||
axises = np.ones(self.ps.shape + (3,))
|
||||
axises[fa + ("xyz".index(plane[0]),)] = 0.0
|
||||
axises[fa + ("xyz".index(plane[1]),)] = 0.0
|
||||
return Quaternions.from_angle_axis(self.ps, axises)
|
||||
|
||||
def directions(self, plane='xz'):
|
||||
dirs = np.zeros((len(self.ps), 3))
|
||||
dirs[..., "xyz".index(plane[0])] = np.sin(self.ps)
|
||||
dirs[..., "xyz".index(plane[1])] = np.cos(self.ps)
|
||||
return dirs
|
||||
|
||||
def normalized(self):
|
||||
xs = np.copy(self.ps)
|
||||
while np.any(xs > np.pi):
|
||||
xs[xs > np.pi] = xs[xs > np.pi] - 2 * np.pi
|
||||
while np.any(xs < -np.pi):
|
||||
xs[xs < -np.pi] = xs[xs < -np.pi] + 2 * np.pi
|
||||
return Pivots(xs)
|
||||
|
||||
def interpolate(self, ws):
|
||||
dir = np.average(self.directions, weights=ws, axis=0)
|
||||
return np.arctan2(dir[2], dir[0])
|
||||
|
||||
def copy(self):
|
||||
return Pivots(np.copy(self.ps))
|
||||
|
||||
@property
|
||||
def shape(self):
|
||||
return self.ps.shape
|
||||
|
||||
@classmethod
|
||||
def from_quaternions(cls, qs, forward='z', plane='xz'):
|
||||
ds = np.zeros(qs.shape + (3,))
|
||||
ds[..., 'xyz'.index(forward)] = 1.0
|
||||
return Pivots.from_directions(qs * ds, plane=plane)
|
||||
|
||||
@classmethod
|
||||
def from_directions(cls, ds, plane='xz'):
|
||||
ys = ds[..., 'xyz'.index(plane[0])]
|
||||
xs = ds[..., 'xyz'.index(plane[1])]
|
||||
return Pivots(np.arctan2(ys, xs))
|
||||
@@ -0,0 +1,481 @@
|
||||
import numpy as np
|
||||
|
||||
|
||||
class Quaternions:
|
||||
"""
|
||||
Quaternions is a wrapper around a numpy ndarray
|
||||
that allows it to act as if it were an narray of
|
||||
a quaternion data type.
|
||||
|
||||
Therefore addition, subtraction, multiplication,
|
||||
division, negation, absolute, are all defined
|
||||
in terms of quaternion operations such as quaternion
|
||||
multiplication.
|
||||
|
||||
This allows for much neater code and many routines
|
||||
which conceptually do the same thing to be written
|
||||
in the same way for point data and for rotation data.
|
||||
|
||||
The Quaternions class has been desgined such that it
|
||||
should support broadcasting and slicing in all of the
|
||||
usual ways.
|
||||
"""
|
||||
|
||||
def __init__(self, qs):
|
||||
if isinstance(qs, np.ndarray):
|
||||
|
||||
if len(qs.shape) == 1:
|
||||
qs = np.array([qs])
|
||||
self.qs = qs
|
||||
return
|
||||
|
||||
if isinstance(qs, Quaternions):
|
||||
self.qs = qs.qs
|
||||
return
|
||||
|
||||
raise TypeError('Quaternions must be constructed from iterable, numpy array, or Quaternions, not %s' % type(qs))
|
||||
|
||||
def __str__(self): return "Quaternions(" + str(self.qs) + ")"
|
||||
def __repr__(self): return "Quaternions(" + repr(self.qs) + ")"
|
||||
""" Helper Methods for Broadcasting and Data extraction """
|
||||
|
||||
@classmethod
|
||||
def _broadcast(cls, sqs, oqs, scalar=False):
|
||||
if isinstance(oqs, float):
|
||||
return sqs, oqs * np.ones(sqs.shape[:-1])
|
||||
|
||||
ss = np.array(sqs.shape) if not scalar else np.array(sqs.shape[:-1])
|
||||
os = np.array(oqs.shape)
|
||||
|
||||
if len(ss) != len(os):
|
||||
raise TypeError('Quaternions cannot broadcast together shapes %s and %s' % (sqs.shape, oqs.shape))
|
||||
|
||||
if np.all(ss == os):
|
||||
return sqs, oqs
|
||||
|
||||
if not np.all((ss == os) | (os == np.ones(len(os))) | (ss == np.ones(len(ss)))):
|
||||
raise TypeError('Quaternions cannot broadcast together shapes %s and %s' % (sqs.shape, oqs.shape))
|
||||
|
||||
sqsn, oqsn = sqs.copy(), oqs.copy()
|
||||
|
||||
for a in np.where(ss == 1)[0]:
|
||||
sqsn = sqsn.repeat(os[a], axis=a)
|
||||
for a in np.where(os == 1)[0]:
|
||||
oqsn = oqsn.repeat(ss[a], axis=a)
|
||||
|
||||
return sqsn, oqsn
|
||||
|
||||
""" Adding Quaterions is just Defined as Multiplication """
|
||||
|
||||
def __add__(self, other): return self * other
|
||||
def __sub__(self, other): return self / other
|
||||
""" Quaterion Multiplication """
|
||||
|
||||
def __mul__(self, other):
|
||||
"""
|
||||
Quaternion multiplication has three main methods.
|
||||
|
||||
When multiplying a Quaternions array by Quaternions
|
||||
normal quaternion multiplication is performed.
|
||||
|
||||
When multiplying a Quaternions array by a vector
|
||||
array of the same shape, where the last axis is 3,
|
||||
it is assumed to be a Quaternion by 3D-Vector
|
||||
multiplication and the 3D-Vectors are rotated
|
||||
in space by the Quaternions.
|
||||
|
||||
When multipplying a Quaternions array by a scalar
|
||||
or vector of different shape it is assumed to be
|
||||
a Quaternions by Scalars multiplication and the
|
||||
Quaternions are scaled using Slerp and the identity
|
||||
quaternions.
|
||||
"""
|
||||
|
||||
""" If Quaternions type do Quaternions * Quaternions """
|
||||
if isinstance(other, Quaternions):
|
||||
sqs, oqs = Quaternions._broadcast(self.qs, other.qs)
|
||||
|
||||
q0 = sqs[..., 0]
|
||||
q1 = sqs[..., 1]
|
||||
q2 = sqs[..., 2]
|
||||
q3 = sqs[..., 3]
|
||||
r0 = oqs[..., 0]
|
||||
r1 = oqs[..., 1]
|
||||
r2 = oqs[..., 2]
|
||||
r3 = oqs[..., 3]
|
||||
|
||||
qs = np.empty(sqs.shape)
|
||||
qs[..., 0] = r0 * q0 - r1 * q1 - r2 * q2 - r3 * q3
|
||||
qs[..., 1] = r0 * q1 + r1 * q0 - r2 * q3 + r3 * q2
|
||||
qs[..., 2] = r0 * q2 + r1 * q3 + r2 * q0 - r3 * q1
|
||||
qs[..., 3] = r0 * q3 - r1 * q2 + r2 * q1 + r3 * q0
|
||||
|
||||
return Quaternions(qs)
|
||||
|
||||
""" If array type do Quaternions * Vectors """
|
||||
if isinstance(other, np.ndarray) and other.shape[-1] == 3:
|
||||
vs = Quaternions(np.concatenate([np.zeros(other.shape[:-1] + (1,)), other], axis=-1))
|
||||
|
||||
return (self * (vs * -self)).imaginaries
|
||||
|
||||
""" If float do Quaternions * Scalars """
|
||||
if isinstance(other, np.ndarray) or isinstance(other, float):
|
||||
return Quaternions.slerp(Quaternions.id_like(self), self, other)
|
||||
|
||||
raise TypeError('Cannot multiply/add Quaternions with type %s' % str(type(other)))
|
||||
|
||||
def __div__(self, other):
|
||||
"""
|
||||
When a Quaternion type is supplied, division is defined
|
||||
as multiplication by the inverse of that Quaternion.
|
||||
|
||||
When a scalar or vector is supplied it is defined
|
||||
as multiplicaion of one over the supplied value.
|
||||
Essentially a scaling.
|
||||
"""
|
||||
|
||||
if isinstance(other, Quaternions):
|
||||
return self * (-other)
|
||||
if isinstance(other, np.ndarray):
|
||||
return self * (1.0 / other)
|
||||
if isinstance(other, float):
|
||||
return self * (1.0 / other)
|
||||
raise TypeError('Cannot divide/subtract Quaternions with type %s' + str(type(other)))
|
||||
|
||||
def __eq__(self, other): return self.qs == other.qs
|
||||
def __ne__(self, other): return self.qs != other.qs
|
||||
|
||||
def __neg__(self):
|
||||
""" Invert Quaternions """
|
||||
return Quaternions(self.qs * np.array([[1, -1, -1, -1]]))
|
||||
|
||||
def __abs__(self):
|
||||
""" Unify Quaternions To Single Pole """
|
||||
qabs = self.normalized().copy()
|
||||
top = np.sum((qabs.qs) * np.array([1, 0, 0, 0]), axis=-1)
|
||||
bot = np.sum((-qabs.qs) * np.array([1, 0, 0, 0]), axis=-1)
|
||||
qabs.qs[top < bot] = -qabs.qs[top < bot]
|
||||
return qabs
|
||||
|
||||
def __iter__(self): return iter(self.qs)
|
||||
def __len__(self): return len(self.qs)
|
||||
|
||||
def __getitem__(self, k): return Quaternions(self.qs[k])
|
||||
def __setitem__(self, k, v): self.qs[k] = v.qs
|
||||
|
||||
@property
|
||||
def lengths(self):
|
||||
return np.sum(self.qs**2.0, axis=-1)**0.5
|
||||
|
||||
@property
|
||||
def reals(self):
|
||||
return self.qs[..., 0]
|
||||
|
||||
@property
|
||||
def imaginaries(self):
|
||||
return self.qs[..., 1:4]
|
||||
|
||||
@property
|
||||
def shape(self): return self.qs.shape[:-1]
|
||||
|
||||
def repeat(self, n, **kwargs):
|
||||
return Quaternions(self.qs.repeat(n, **kwargs))
|
||||
|
||||
def normalized(self):
|
||||
return Quaternions(self.qs / self.lengths[..., np.newaxis])
|
||||
|
||||
def log(self):
|
||||
norm = abs(self.normalized())
|
||||
imgs = norm.imaginaries
|
||||
lens = np.sqrt(np.sum(imgs**2, axis=-1))
|
||||
lens = np.arctan2(lens, norm.reals) / (lens + 1e-10)
|
||||
return imgs * lens[..., np.newaxis]
|
||||
|
||||
def constrained(self, axis):
|
||||
|
||||
rl = self.reals
|
||||
im = np.sum(axis * self.imaginaries, axis=-1)
|
||||
|
||||
t1 = -2 * np.arctan2(rl, im) + np.pi
|
||||
t2 = -2 * np.arctan2(rl, im) - np.pi
|
||||
|
||||
top = Quaternions.exp(axis[np.newaxis] * (t1[:, np.newaxis] / 2.0))
|
||||
bot = Quaternions.exp(axis[np.newaxis] * (t2[:, np.newaxis] / 2.0))
|
||||
img = self.dot(top) > self.dot(bot)
|
||||
|
||||
ret = top.copy()
|
||||
ret[img] = top[img]
|
||||
ret[~img] = bot[~img]
|
||||
return ret
|
||||
|
||||
def constrained_x(self): return self.constrained(np.array([1, 0, 0]))
|
||||
def constrained_y(self): return self.constrained(np.array([0, 1, 0]))
|
||||
def constrained_z(self): return self.constrained(np.array([0, 0, 1]))
|
||||
|
||||
def dot(self, q): return np.sum(self.qs * q.qs, axis=-1)
|
||||
|
||||
def copy(self): return Quaternions(np.copy(self.qs))
|
||||
|
||||
def reshape(self, s):
|
||||
self.qs.reshape(s)
|
||||
return self
|
||||
|
||||
def interpolate(self, ws):
|
||||
return Quaternions.exp(np.average(abs(self).log, axis=0, weights=ws))
|
||||
|
||||
def euler(self, order='xyz'):
|
||||
|
||||
q = self.normalized().qs
|
||||
q0 = q[..., 0]
|
||||
q1 = q[..., 1]
|
||||
q2 = q[..., 2]
|
||||
q3 = q[..., 3]
|
||||
es = np.zeros(self.shape + (3,))
|
||||
|
||||
if order == 'xyz':
|
||||
es[..., 0] = np.arctan2(2 * (q0 * q1 + q2 * q3), 1 - 2 * (q1 * q1 + q2 * q2))
|
||||
es[..., 1] = np.arcsin((2 * (q0 * q2 - q3 * q1)).clip(-1, 1))
|
||||
es[..., 2] = np.arctan2(2 * (q0 * q3 + q1 * q2), 1 - 2 * (q2 * q2 + q3 * q3))
|
||||
elif order == 'yzx':
|
||||
es[..., 0] = np.arctan2(2 * (q1 * q0 - q2 * q3), -q1 * q1 + q2 * q2 - q3 * q3 + q0 * q0)
|
||||
es[..., 1] = np.arctan2(2 * (q2 * q0 - q1 * q3), q1 * q1 - q2 * q2 - q3 * q3 + q0 * q0)
|
||||
es[..., 2] = np.arcsin((2 * (q1 * q2 + q3 * q0)).clip(-1, 1))
|
||||
else:
|
||||
raise NotImplementedError('Cannot convert from ordering %s' % order)
|
||||
|
||||
"""
|
||||
|
||||
# These conversion don't appear to work correctly for Maya.
|
||||
# http://bediyap.com/programming/convert-quaternion-to-euler-rotations/
|
||||
|
||||
if order == 'xyz':
|
||||
es[fa + (0,)] = np.arctan2(2 * (q0 * q3 - q1 * q2), q0 * q0 + q1 * q1 - q2 * q2 - q3 * q3)
|
||||
es[fa + (1,)] = np.arcsin((2 * (q1 * q3 + q0 * q2)).clip(-1,1))
|
||||
es[fa + (2,)] = np.arctan2(2 * (q0 * q1 - q2 * q3), q0 * q0 - q1 * q1 - q2 * q2 + q3 * q3)
|
||||
elif order == 'yzx':
|
||||
es[fa + (0,)] = np.arctan2(2 * (q0 * q1 - q2 * q3), q0 * q0 - q1 * q1 + q2 * q2 - q3 * q3)
|
||||
es[fa + (1,)] = np.arcsin((2 * (q1 * q2 + q0 * q3)).clip(-1,1))
|
||||
es[fa + (2,)] = np.arctan2(2 * (q0 * q2 - q1 * q3), q0 * q0 + q1 * q1 - q2 * q2 - q3 * q3)
|
||||
elif order == 'zxy':
|
||||
es[fa + (0,)] = np.arctan2(2 * (q0 * q2 - q1 * q3), q0 * q0 - q1 * q1 - q2 * q2 + q3 * q3)
|
||||
es[fa + (1,)] = np.arcsin((2 * (q0 * q1 + q2 * q3)).clip(-1,1))
|
||||
es[fa + (2,)] = np.arctan2(2 * (q0 * q3 - q1 * q2), q0 * q0 - q1 * q1 + q2 * q2 - q3 * q3)
|
||||
elif order == 'xzy':
|
||||
es[fa + (0,)] = np.arctan2(2 * (q0 * q2 + q1 * q3), q0 * q0 + q1 * q1 - q2 * q2 - q3 * q3)
|
||||
es[fa + (1,)] = np.arcsin((2 * (q0 * q3 - q1 * q2)).clip(-1,1))
|
||||
es[fa + (2,)] = np.arctan2(2 * (q0 * q1 + q2 * q3), q0 * q0 - q1 * q1 + q2 * q2 - q3 * q3)
|
||||
elif order == 'yxz':
|
||||
es[fa + (0,)] = np.arctan2(2 * (q1 * q2 + q0 * q3), q0 * q0 - q1 * q1 + q2 * q2 - q3 * q3)
|
||||
es[fa + (1,)] = np.arcsin((2 * (q0 * q1 - q2 * q3)).clip(-1,1))
|
||||
es[fa + (2,)] = np.arctan2(2 * (q1 * q3 + q0 * q2), q0 * q0 - q1 * q1 - q2 * q2 + q3 * q3)
|
||||
elif order == 'zyx':
|
||||
es[fa + (0,)] = np.arctan2(2 * (q0 * q1 + q2 * q3), q0 * q0 - q1 * q1 - q2 * q2 + q3 * q3)
|
||||
es[fa + (1,)] = np.arcsin((2 * (q0 * q2 - q1 * q3)).clip(-1,1))
|
||||
es[fa + (2,)] = np.arctan2(2 * (q0 * q3 + q1 * q2), q0 * q0 + q1 * q1 - q2 * q2 - q3 * q3)
|
||||
else:
|
||||
raise KeyError('Unknown ordering %s' % order)
|
||||
|
||||
"""
|
||||
|
||||
# https://github.com/ehsan/ogre/blob/master/OgreMain/src/OgreMatrix3.cpp
|
||||
# Use this class and convert from matrix
|
||||
|
||||
return es
|
||||
|
||||
def average(self):
|
||||
|
||||
if len(self.shape) == 1:
|
||||
|
||||
import numpy.core.umath_tests as ut
|
||||
system = ut.matrix_multiply(self.qs[:, :, np.newaxis], self.qs[:, np.newaxis, :]).sum(axis=0)
|
||||
w, v = np.linalg.eigh(system)
|
||||
qiT_dot_qref = (self.qs[:, :, np.newaxis] * v[np.newaxis, :, :]).sum(axis=1)
|
||||
return Quaternions(v[:, np.argmin((1.-qiT_dot_qref**2).sum(axis=0))])
|
||||
|
||||
else:
|
||||
|
||||
raise NotImplementedError('Cannot average multi-dimensionsal Quaternions')
|
||||
|
||||
def angle_axis(self):
|
||||
|
||||
norm = self.normalized()
|
||||
s = np.sqrt(1 - (norm.reals**2.0))
|
||||
s[s == 0] = 0.001
|
||||
|
||||
angles = 2.0 * np.arccos(norm.reals)
|
||||
axis = norm.imaginaries / s[..., np.newaxis]
|
||||
|
||||
return angles, axis
|
||||
|
||||
def transforms(self):
|
||||
|
||||
qw = self.qs[..., 0]
|
||||
qx = self.qs[..., 1]
|
||||
qy = self.qs[..., 2]
|
||||
qz = self.qs[..., 3]
|
||||
|
||||
x2 = qx + qx
|
||||
y2 = qy + qy
|
||||
z2 = qz + qz
|
||||
xx = qx * x2
|
||||
yy = qy * y2
|
||||
wx = qw * x2
|
||||
xy = qx * y2
|
||||
yz = qy * z2
|
||||
wy = qw * y2
|
||||
xz = qx * z2
|
||||
zz = qz * z2
|
||||
wz = qw * z2
|
||||
|
||||
m = np.empty(self.shape + (3, 3))
|
||||
m[..., 0, 0] = 1.0 - (yy + zz)
|
||||
m[..., 0, 1] = xy - wz
|
||||
m[..., 0, 2] = xz + wy
|
||||
m[..., 1, 0] = xy + wz
|
||||
m[..., 1, 1] = 1.0 - (xx + zz)
|
||||
m[..., 1, 2] = yz - wx
|
||||
m[..., 2, 0] = xz - wy
|
||||
m[..., 2, 1] = yz + wx
|
||||
m[..., 2, 2] = 1.0 - (xx + yy)
|
||||
|
||||
return m
|
||||
|
||||
def ravel(self):
|
||||
return self.qs.ravel()
|
||||
|
||||
@classmethod
|
||||
def id(cls, n):
|
||||
|
||||
if isinstance(n, tuple):
|
||||
qs = np.zeros(n + (4,))
|
||||
qs[..., 0] = 1.0
|
||||
return Quaternions(qs)
|
||||
|
||||
if isinstance(n, int):
|
||||
qs = np.zeros((n, 4))
|
||||
qs[:, 0] = 1.0
|
||||
return Quaternions(qs)
|
||||
|
||||
raise TypeError('Cannot Construct Quaternion from %s type' % str(type(n)))
|
||||
|
||||
@classmethod
|
||||
def id_like(cls, a):
|
||||
qs = np.zeros(a.shape + (4,))
|
||||
qs[..., 0] = 1.0
|
||||
return Quaternions(qs)
|
||||
|
||||
@classmethod
|
||||
def exp(cls, ws):
|
||||
|
||||
ts = np.sum(ws**2.0, axis=-1)**0.5
|
||||
ts[ts == 0] = 0.001
|
||||
ls = np.sin(ts) / ts
|
||||
|
||||
qs = np.empty(ws.shape[:-1] + (4,))
|
||||
qs[..., 0] = np.cos(ts)
|
||||
qs[..., 1] = ws[..., 0] * ls
|
||||
qs[..., 2] = ws[..., 1] * ls
|
||||
qs[..., 3] = ws[..., 2] * ls
|
||||
|
||||
return Quaternions(qs).normalized()
|
||||
|
||||
@classmethod
|
||||
def slerp(cls, q0s, q1s, a):
|
||||
|
||||
fst, snd = cls._broadcast(q0s.qs, q1s.qs)
|
||||
fst, a = cls._broadcast(fst, a, scalar=True)
|
||||
snd, a = cls._broadcast(snd, a, scalar=True)
|
||||
|
||||
len = np.sum(fst * snd, axis=-1)
|
||||
|
||||
neg = len < 0.0
|
||||
len[neg] = -len[neg]
|
||||
snd[neg] = -snd[neg]
|
||||
|
||||
amount0 = np.zeros(a.shape)
|
||||
amount1 = np.zeros(a.shape)
|
||||
|
||||
linear = (1.0 - len) < 0.01
|
||||
omegas = np.arccos(len[~linear])
|
||||
sinoms = np.sin(omegas)
|
||||
|
||||
amount0[linear] = 1.0 - a[linear]
|
||||
amount1[linear] = a[linear]
|
||||
amount0[~linear] = np.sin((1.0 - a[~linear]) * omegas) / sinoms
|
||||
amount1[~linear] = np.sin(a[~linear] * omegas) / sinoms
|
||||
|
||||
return Quaternions(
|
||||
amount0[..., np.newaxis] * fst +
|
||||
amount1[..., np.newaxis] * snd)
|
||||
|
||||
@classmethod
|
||||
def between(cls, v0s, v1s):
|
||||
a = np.cross(v0s, v1s)
|
||||
w = np.sqrt((v0s**2).sum(axis=-1) * (v1s**2).sum(axis=-1)) + (v0s * v1s).sum(axis=-1)
|
||||
return Quaternions(np.concatenate([w[..., np.newaxis], a], axis=-1)).normalized()
|
||||
|
||||
@classmethod
|
||||
def from_angle_axis(cls, angles, axis):
|
||||
axis = axis / (np.sqrt(np.sum(axis**2, axis=-1)) + 1e-10)[..., np.newaxis]
|
||||
sines = np.sin(angles / 2.0)[..., np.newaxis]
|
||||
cosines = np.cos(angles / 2.0)[..., np.newaxis]
|
||||
return Quaternions(np.concatenate([cosines, axis * sines], axis=-1))
|
||||
|
||||
@classmethod
|
||||
def from_euler(cls, es, order='xyz', world=False):
|
||||
|
||||
axis = {
|
||||
'x': np.array([1, 0, 0]),
|
||||
'y': np.array([0, 1, 0]),
|
||||
'z': np.array([0, 0, 1]),
|
||||
}
|
||||
|
||||
q0s = Quaternions.from_angle_axis(es[..., 0], axis[order[0]])
|
||||
q1s = Quaternions.from_angle_axis(es[..., 1], axis[order[1]])
|
||||
q2s = Quaternions.from_angle_axis(es[..., 2], axis[order[2]])
|
||||
|
||||
return (q2s * (q1s * q0s)) if world else (q0s * (q1s * q2s))
|
||||
|
||||
@classmethod
|
||||
def from_transforms(cls, ts):
|
||||
|
||||
d0, d1, d2 = ts[..., 0, 0], ts[..., 1, 1], ts[..., 2, 2]
|
||||
|
||||
q0 = (d0 + d1 + d2 + 1.0) / 4.0
|
||||
q1 = (d0 - d1 - d2 + 1.0) / 4.0
|
||||
q2 = (-d0 + d1 - d2 + 1.0) / 4.0
|
||||
q3 = (-d0 - d1 + d2 + 1.0) / 4.0
|
||||
|
||||
q0 = np.sqrt(q0.clip(0, None))
|
||||
q1 = np.sqrt(q1.clip(0, None))
|
||||
q2 = np.sqrt(q2.clip(0, None))
|
||||
q3 = np.sqrt(q3.clip(0, None))
|
||||
|
||||
c0 = (q0 >= q1) & (q0 >= q2) & (q0 >= q3)
|
||||
c1 = (q1 >= q0) & (q1 >= q2) & (q1 >= q3)
|
||||
c2 = (q2 >= q0) & (q2 >= q1) & (q2 >= q3)
|
||||
c3 = (q3 >= q0) & (q3 >= q1) & (q3 >= q2)
|
||||
|
||||
q1[c0] *= np.sign(ts[c0, 2, 1] - ts[c0, 1, 2])
|
||||
q2[c0] *= np.sign(ts[c0, 0, 2] - ts[c0, 2, 0])
|
||||
q3[c0] *= np.sign(ts[c0, 1, 0] - ts[c0, 0, 1])
|
||||
|
||||
q0[c1] *= np.sign(ts[c1, 2, 1] - ts[c1, 1, 2])
|
||||
q2[c1] *= np.sign(ts[c1, 1, 0] + ts[c1, 0, 1])
|
||||
q3[c1] *= np.sign(ts[c1, 0, 2] + ts[c1, 2, 0])
|
||||
|
||||
q0[c2] *= np.sign(ts[c2, 0, 2] - ts[c2, 2, 0])
|
||||
q1[c2] *= np.sign(ts[c2, 1, 0] + ts[c2, 0, 1])
|
||||
q3[c2] *= np.sign(ts[c2, 2, 1] + ts[c2, 1, 2])
|
||||
|
||||
q0[c3] *= np.sign(ts[c3, 1, 0] - ts[c3, 0, 1])
|
||||
q1[c3] *= np.sign(ts[c3, 2, 0] + ts[c3, 0, 2])
|
||||
q2[c3] *= np.sign(ts[c3, 2, 1] + ts[c3, 1, 2])
|
||||
|
||||
qs = np.empty(ts.shape[:-2] + (4,))
|
||||
qs[..., 0] = q0
|
||||
qs[..., 1] = q1
|
||||
qs[..., 2] = q2
|
||||
qs[..., 3] = q3
|
||||
|
||||
return cls(qs)
|
||||
@@ -0,0 +1,4 @@
|
||||
import sys
|
||||
import os
|
||||
BASEPATH = os.path.dirname(__file__)
|
||||
sys.path.insert(0, BASEPATH)
|
||||
+60
@@ -0,0 +1,60 @@
|
||||
smpl_body_model: None # _DATA/data/mano
|
||||
|
||||
offsets:
|
||||
right: [[ 0.0957, 0.0064, 0.0062],
|
||||
[ 0.0076, 0.0012, 0.0269],
|
||||
[-0.0251, 0.0052, 0.0291],
|
||||
[-0.0473, 0.0039, 0.0290],
|
||||
[ 0.0010, 0.0049, 0.0028],
|
||||
[-0.0302, 0.0068, -0.0028],
|
||||
[-0.0531, 0.0055, -0.0067],
|
||||
[ 0.0269, -0.0036, -0.0370],
|
||||
[ 0.0099, -0.0035, -0.0495],
|
||||
[-0.0060, -0.0042, -0.0599],
|
||||
[ 0.0139, 0.0024, -0.0205],
|
||||
[-0.0144, 0.0045, -0.0256],
|
||||
[-0.0379, 0.0028, -0.0332],
|
||||
[ 0.0716, -0.0091, 0.0320],
|
||||
[ 0.0519, -0.0082, 0.0557],
|
||||
[ 0.0297, -0.0137, 0.0702]]
|
||||
|
||||
|
||||
left: [
|
||||
[-0.000876, 0.909315, 0.027821],
|
||||
[ 0.071361, -0.089584, -0.008046],
|
||||
[ 0.030669, -0.364209, -0.006689],
|
||||
[-0.011554, -0.383348, -0.043502],
|
||||
[ 0.023338, -0.054645, 0.114370],
|
||||
[-0.069012, -0.088960, -0.004796],
|
||||
[-0.036152, -0.370650, -0.009185],
|
||||
[ 0.014029, -0.383638, -0.041892],
|
||||
[-0.022043, -0.046410, 0.117900],
|
||||
[-0.002508, 0.103257, -0.022185],
|
||||
[ 0.003581, 0.127658, -0.001713],
|
||||
[ 0.002027, 0.049072, 0.027867],
|
||||
[-0.001963, 0.208243, -0.049765],
|
||||
[ 0.003517, 0.062322, 0.050205],
|
||||
[ 0.075298, 0.117780, -0.036875],
|
||||
[ 0.085317, 0.031739, -0.007293],
|
||||
[ 0.251247, -0.011967, -0.027518],
|
||||
[ 0.238019, 0.009007, 0.000044],
|
||||
[ 0.079668, -0.009683, -0.013206],
|
||||
[-0.077033, 0.115606, -0.041811],
|
||||
[-0.089203, 0.032785, -0.009802],
|
||||
[-0.245990, -0.013152, -0.020162],
|
||||
[-0.245177, 0.008622, -0.003532],
|
||||
[-0.080400, -0.007248, -0.010419]
|
||||
]
|
||||
|
||||
parents: [-1, 0, 1, 2, 0, 4, 5, 0, 7, 8, 0, 10, 11, 0, 13, 14]
|
||||
|
||||
joint_names: ['wrist', 'index1', 'index2', 'index3', 'middle1', 'middle2', 'middle3', 'pinky1', 'pinky2', 'pinky3',
|
||||
'ring1', 'ring2', 'ring3', 'thumb1', 'thumb2', 'thumb3']
|
||||
|
||||
|
||||
joints_to_use: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]
|
||||
|
||||
leaf_joints: [16, 17, 18, 19, 20]
|
||||
|
||||
wrist_index: [0]
|
||||
|
||||
@@ -0,0 +1,165 @@
|
||||
import math
|
||||
import os
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
import torch
|
||||
import torch.nn.init
|
||||
import torch.optim
|
||||
|
||||
|
||||
def weights_init(init_type='default'):
|
||||
def init_fun(m):
|
||||
classname = m.__class__.__name__
|
||||
if (classname.find('Conv') == 0 or classname.find('Linear') == 0) and hasattr(m, 'weight'):
|
||||
if init_type == 'gaussian':
|
||||
torch.nn.init.normal_(m.weight.data)
|
||||
elif init_type == 'xavier':
|
||||
torch.nn.init.xavier_normal_(m.weight.data, gain=math.sqrt(2))
|
||||
elif init_type == 'kaiming':
|
||||
torch.nn.init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
|
||||
elif init_type == 'orthogonal':
|
||||
torch.nn.init.orthogonal_(m.weight.data, gain=math.sqrt(2))
|
||||
elif init_type == 'default':
|
||||
torch.nn.init.normal_(m.weight.data)
|
||||
weight_norm = m.weight.pow(2).sum(2, keepdim=True).sum(1, keepdim=True).add(1e-8).sqrt()
|
||||
m.weight.data.div_(weight_norm)
|
||||
else:
|
||||
raise NotImplementedError(f"Unsupported initialization: {init_type}")
|
||||
if hasattr(m, 'bias') and m.bias is not None:
|
||||
torch.nn.init.constant_(m.bias.data, 0.0)
|
||||
return init_fun
|
||||
|
||||
|
||||
class BaseModel(ABC):
|
||||
"""This class is an abstract base class (ABC) for models.
|
||||
To create a subclass, you need to implement the following five functions:
|
||||
-- <__init__>: initialize the class; first call BaseModel.__init__(self, opt).
|
||||
-- <set_input>: unpack data from dataset and apply preprocessing.
|
||||
-- <forward>: produce intermediate results.
|
||||
-- <optimize_parameters>: calculate losses, gradients, and update network weights.
|
||||
"""
|
||||
|
||||
def __init__(self, args):
|
||||
self.args = args
|
||||
self.is_train = args.is_train
|
||||
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||||
self.model_save_dir = os.path.join(args.save_dir, 'checkpoints') # save all the checkpoints to save_dir
|
||||
|
||||
if self.is_train and args.log:
|
||||
import time
|
||||
from tensorboardX import SummaryWriter
|
||||
from .loss_record import LossRecorder
|
||||
|
||||
log_dir = os.path.join(args.save_dir, 'logs')
|
||||
if args.epoch_begin != 0:
|
||||
all = [d for d in os.listdir(log_dir)]
|
||||
if len(all) == 0:
|
||||
raise RuntimeError(f'Empty logging path {log_dir}')
|
||||
timestamp = sorted(all)[-1]
|
||||
else:
|
||||
timestamp = time.strftime("%Y-%m-%d-%H:%M:%S", time.localtime(time.time()))
|
||||
self.log_path = os.path.join(log_dir, timestamp)
|
||||
os.makedirs(self.log_path, exist_ok=True)
|
||||
self.writer = SummaryWriter(self.log_path)
|
||||
self.loss_recorder = LossRecorder(self.writer)
|
||||
|
||||
self.epoch_cnt = 0
|
||||
self.schedulers = []
|
||||
self.optimizers = []
|
||||
self.models = []
|
||||
|
||||
@abstractmethod
|
||||
def set_input(self, input):
|
||||
"""Unpack input data from the dataloader and perform necessary pre-processing steps.
|
||||
Parameters:
|
||||
input (dict): includes the data itself and its metadata information.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def compute_test_result(self):
|
||||
"""
|
||||
After forward, do something like output bvh, get error value
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def forward(self):
|
||||
"""Run forward pass; called by both functions <optimize_parameters> and <test>."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def optimize_parameters(self):
|
||||
"""Calculate losses, gradients, and update network weights; called in every training iteration"""
|
||||
pass
|
||||
|
||||
def get_scheduler(self, optimizer):
|
||||
if self.args.scheduler.name == 'linear':
|
||||
def lambda_rule(epoch):
|
||||
lr_l = 1.0 - max(0, epoch - self.args.n_epochs_origin) / float(self.args.n_epochs_decay + 1)
|
||||
return lr_l
|
||||
return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_rule)
|
||||
if self.args.scheduler.name == 'StepLR':
|
||||
print('StepLR scheduler set')
|
||||
return torch.optim.lr_scheduler.StepLR(optimizer, self.args.scheduler.step_size, self.args.scheduler.gamma, verbose=self.args.verbose)
|
||||
if self.args.scheduler.name == 'Plateau':
|
||||
print('ReduceLROnPlateau shceduler set')
|
||||
return torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=self.args.scheduler.factor, threshold=self.args.scheduler.threshold, patience=5, verbose=True)
|
||||
if self.args.scheduler.name == 'MultiStep':
|
||||
print('MultiStep shceduler set')
|
||||
return torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=self.args.scheduler.milestones, gamma=self.args.scheduler.gamma, verbose=self.args.verbose)
|
||||
else:
|
||||
print('No scheduler set')
|
||||
return None
|
||||
|
||||
def setup(self):
|
||||
"""Load and print networks; create schedulers
|
||||
Parameters:
|
||||
opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
|
||||
"""
|
||||
if self.is_train:
|
||||
self.schedulers = [self.get_scheduler(optimizer) for optimizer in self.optimizers]
|
||||
|
||||
def epoch(self):
|
||||
if self.args.verbose:
|
||||
self.loss_recorder.epoch()
|
||||
for scheduler in self.schedulers:
|
||||
if scheduler is not None:
|
||||
if self.args.scheduler.name == 'Plateau':
|
||||
loss = self.loss_recorder.losses['total_loss'].current()
|
||||
scheduler.step(loss[1])
|
||||
else:
|
||||
scheduler.step()
|
||||
self.epoch_cnt += 1
|
||||
|
||||
def test(self):
|
||||
"""Forward function used in test time.
|
||||
This function wraps <forward> function in no_grad() so we don't save intermediate steps for backprop
|
||||
It also calls <compute_visuals> to produce additional visualization results
|
||||
"""
|
||||
with torch.no_grad():
|
||||
self.forward()
|
||||
self.compute_test_result()
|
||||
|
||||
def train(self):
|
||||
for model in self.models:
|
||||
model.train()
|
||||
for param in model.parameters():
|
||||
param.requires_grad = True
|
||||
|
||||
def eval(self):
|
||||
for model in self.models:
|
||||
model.eval()
|
||||
for param in model.parameters():
|
||||
param.requires_grad = False
|
||||
|
||||
def print(self):
|
||||
for model in self.models:
|
||||
print(model)
|
||||
|
||||
def count_params(self):
|
||||
params = 0
|
||||
for model in self.models:
|
||||
params += sum(p.numel() for p in model.parameters() if p.requires_grad)
|
||||
|
||||
return params
|
||||
+291
@@ -0,0 +1,291 @@
|
||||
import os
|
||||
|
||||
import holden.BVH as BVH
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from holden.Animation import Animation
|
||||
from holden.Quaternions import Quaternions
|
||||
from human_body_prior.tools.omni_tools import copy2cpu as c2c
|
||||
from rotations import matrix_to_quaternion, quaternion_to_axis_angle, rotation_6d_to_matrix
|
||||
from utils import align_joints, compute_trajectory
|
||||
|
||||
from .base_model import BaseModel
|
||||
from .fk import ForwardKinematicsLayer
|
||||
from .losses import GeodesicLoss
|
||||
from .neural_motion import NeuralMotionField
|
||||
|
||||
|
||||
class Architecture(BaseModel):
|
||||
def __init__(self, args, ngpu):
|
||||
super(Architecture, self).__init__(args)
|
||||
|
||||
self.args = args
|
||||
|
||||
if args.amass_data:
|
||||
self.fk = ForwardKinematicsLayer(args)
|
||||
else:
|
||||
self.fk = ForwardKinematicsLayer(parents=args.animal.parents, positions=args.animal.offsets)
|
||||
|
||||
self.field = NeuralMotionField(args.nemf).to(self.device)
|
||||
if args.multi_gpu is True:
|
||||
self.field = nn.DataParallel(self.field, device_ids=range(ngpu))
|
||||
|
||||
self.models = [self.field]
|
||||
|
||||
ordermap = {
|
||||
'x': 0,
|
||||
'y': 1,
|
||||
'z': 2,
|
||||
}
|
||||
self.up_axis = ordermap[self.args.data.up]
|
||||
self.v_axis = [x for x in ordermap.values() if x != self.up_axis]
|
||||
|
||||
print(self.up_axis)
|
||||
print(self.v_axis)
|
||||
|
||||
self.input_data = dict()
|
||||
self.recon_data = dict()
|
||||
|
||||
if self.is_train:
|
||||
self.optimizer = torch.optim.Adam(self.field.parameters(), args.learning_rate)
|
||||
self.optimizers = [self.optimizer]
|
||||
self.criterion_rec = nn.L1Loss() if args.l1_loss else nn.MSELoss()
|
||||
|
||||
self.criterion_geo = GeodesicLoss().to(self.device)
|
||||
self.fps = args.data.fps
|
||||
if args.bvh_viz:
|
||||
self.viz_dir = os.path.join(args.save_dir, 'results', 'bvh')
|
||||
else:
|
||||
self.viz_dir = os.path.join(args.save_dir, 'results', 'smpl')
|
||||
|
||||
def set_input(self, input):
|
||||
self.input_data = {k: v.float().to(self.device) for k, v in input.items() if k in ['pos', 'global_xform', 'root_orient', 'root_vel', 'trans']}
|
||||
self.input_data['rotmat'] = rotation_6d_to_matrix(self.input_data['global_xform'])
|
||||
|
||||
def forward(self, step=1):
|
||||
b_size, t_length, n_joints = self.input_data['pos'].shape[:3]
|
||||
|
||||
t = torch.arange(start=0, end=t_length, step=step).unsqueeze(0) # (1, T)
|
||||
t = t / t_length * 2 - 1 # map t to [-1, 1]
|
||||
t = t.expand(b_size, -1).unsqueeze(-1).to(self.device) # (B, T, 1)
|
||||
f, _ = self.field(t) # (B, T, D), rot6d + root orient + root vel + root height
|
||||
# import ipdb; ipdb.set_trace()
|
||||
self.recon_data = dict()
|
||||
rot6d_recon = f[:, :, :n_joints * 6] # (B, T, J x 6)
|
||||
rot6d_recon = rot6d_recon.view(b_size, -1, n_joints, 6) # (B, T, J, 6)
|
||||
rotmat_recon = rotation_6d_to_matrix(rot6d_recon) # (B, T, J, 3, 3)
|
||||
if self.args.data.root_transform:
|
||||
identity = torch.zeros(b_size, t.shape[1], 3, 3).to(self.device)
|
||||
identity[:, :, 0, 0] = 1
|
||||
identity[:, :, 1, 1] = 1
|
||||
identity[:, :, 2, 2] = 1
|
||||
rotmat_recon[:, :, 0] = identity
|
||||
# rotmat_recon[:, :, -2] = rotmat_recon[:, :, -4]
|
||||
# rotmat_recon[:, :, -1] = rotmat_recon[:, :, -3]
|
||||
local_rotmat = self.fk.global_to_local(rotmat_recon.view(-1, n_joints, 3, 3)) # (B x T, J, 3. 3)
|
||||
pos_recon, _ = self.fk(local_rotmat) # (B x T, J, 3)
|
||||
pos_recon = pos_recon.contiguous().view(b_size, -1, n_joints, 3) # (B, T, J, 3)
|
||||
self.recon_data['rotmat'] = rotmat_recon
|
||||
self.recon_data['pos'] = pos_recon
|
||||
|
||||
if self.args.data.root_transform:
|
||||
root_orient = f[:, :, n_joints * 6:n_joints * 6 + 6] # (B, T, 6)
|
||||
self.recon_data['root_orient'] = root_orient
|
||||
else:
|
||||
self.recon_data['root_orient'] = rot6d_recon[:, :, 0] # (B, T, 6)
|
||||
|
||||
root_vel = f[:, :, -4:-1] # (B, T, 3)
|
||||
root_height = f[:, :, -1] # (B, T)
|
||||
|
||||
self.recon_data['root_vel'] = root_vel
|
||||
self.recon_data['root_height'] = root_height
|
||||
|
||||
dt = 1.0 / self.args.data.fps * step
|
||||
origin = torch.zeros(b_size, 3).to(self.device)
|
||||
trans = compute_trajectory(root_vel, root_height, origin, dt, up_axis=self.args.data.up)
|
||||
self.recon_data['trans'] = trans
|
||||
|
||||
def backward(self, validation=False):
|
||||
loss = 0
|
||||
|
||||
root_orient = rotation_6d_to_matrix(self.recon_data['root_orient'])
|
||||
root_orient_gt = rotation_6d_to_matrix(self.input_data['root_orient'])
|
||||
if self.args.geodesic_loss:
|
||||
orient_recon_loss = self.criterion_geo(root_orient.view(-1, 3, 3), root_orient_gt.view(-1, 3, 3))
|
||||
else:
|
||||
orient_recon_loss = self.criterion_rec(root_orient, root_orient_gt)
|
||||
self.loss_recorder.add_scalar('orient_recon_loss', orient_recon_loss, validation=validation)
|
||||
loss += self.args.lambda_orient * orient_recon_loss
|
||||
|
||||
if self.args.geodesic_loss:
|
||||
rotmat_recon_loss = self.criterion_geo(self.recon_data['rotmat'].view(-1, 3, 3), self.input_data['rotmat'].view(-1, 3, 3))
|
||||
else:
|
||||
rotmat_recon_loss = self.criterion_rec(self.recon_data['rotmat'], self.input_data['rotmat'])
|
||||
self.loss_recorder.add_scalar('rotmat_recon_loss', rotmat_recon_loss, validation=validation)
|
||||
loss += self.args.lambda_rotmat * rotmat_recon_loss
|
||||
|
||||
pos_recon_loss = self.criterion_rec(self.recon_data['pos'], self.input_data['pos'])
|
||||
self.loss_recorder.add_scalar('pos_recon_loss', pos_recon_loss, validation=validation)
|
||||
loss += self.args.lambda_pos * pos_recon_loss
|
||||
|
||||
v_loss = self.criterion_rec(self.recon_data['root_vel'], self.input_data['root_vel'])
|
||||
up_loss = self.criterion_rec(self.recon_data['root_height'], self.input_data['trans'][..., 'xyz'.index(self.args.data.up)])
|
||||
self.loss_recorder.add_scalar('v_loss', v_loss, validation=validation)
|
||||
self.loss_recorder.add_scalar('up_loss', up_loss, validation=validation)
|
||||
loss += self.args.lambda_v * v_loss + self.args.lambda_up * up_loss
|
||||
|
||||
origin = self.input_data['trans'][:, 0]
|
||||
trans = self.recon_data['trans']
|
||||
trans[..., self.v_axis] = trans[..., self.v_axis] + origin[..., self.v_axis].unsqueeze(1)
|
||||
trans_loss = self.criterion_rec(trans, self.input_data['trans'])
|
||||
self.loss_recorder.add_scalar('trans_loss', trans_loss, validation=validation)
|
||||
loss += self.args.lambda_trans * trans_loss
|
||||
|
||||
self.loss_recorder.add_scalar('total_loss', loss, validation=validation)
|
||||
|
||||
if not validation:
|
||||
loss.backward()
|
||||
|
||||
def optimize_parameters(self):
|
||||
self.optimizer.zero_grad()
|
||||
self.forward()
|
||||
self.backward(validation=False)
|
||||
self.optimizer.step()
|
||||
|
||||
def report_errors(self):
|
||||
root_orient = rotation_6d_to_matrix(self.recon_data['root_orient']) # (B, T, 3, 3)
|
||||
root_orient_gt = rotation_6d_to_matrix(self.input_data['root_orient']) # (B, T, 3, 3)
|
||||
orientation_error = self.criterion_geo(root_orient.view(-1, 3, 3), root_orient_gt.view(-1, 3, 3))
|
||||
|
||||
local_rotmat = self.fk.global_to_local(self.recon_data['rotmat'].view(-1, self.args.smpl.joint_num, 3, 3)) # (B x T, J, 3, 3)
|
||||
local_rotmat_gt = self.fk.global_to_local(self.input_data['rotmat'].view(-1, self.args.smpl.joint_num, 3, 3)) # (B x T, J, 3, 3)
|
||||
rotation_error = self.criterion_geo(local_rotmat[:, 1:].reshape(-1, 3, 3), local_rotmat_gt[:, 1:].reshape(-1, 3, 3))
|
||||
|
||||
pos = self.recon_data['pos'] # (B, T, J, 3)
|
||||
pos_gt = self.input_data['pos'] # (B, T, J, 3)
|
||||
position_error = torch.linalg.norm((pos - pos_gt), dim=-1).mean()
|
||||
|
||||
origin = self.input_data['trans'][:, 0]
|
||||
trans = self.recon_data['trans']
|
||||
trans[..., self.v_axis] = trans[..., self.v_axis] + origin[..., self.v_axis].unsqueeze(1)
|
||||
trans_gt = self.input_data['trans'] # (B, T, 3)
|
||||
translation_error = torch.linalg.norm((trans - trans_gt), dim=-1).mean()
|
||||
|
||||
return {
|
||||
'rotation': c2c(rotation_error),
|
||||
'position': c2c(position_error),
|
||||
'orientation': c2c(orientation_error),
|
||||
'translation': c2c(translation_error)
|
||||
}
|
||||
|
||||
def verbose(self):
|
||||
res = {}
|
||||
for loss in self.loss_recorder.losses.values():
|
||||
res[loss.name] = {'train': loss.current()[0]}
|
||||
|
||||
return res
|
||||
|
||||
def save(self, optimal=False):
|
||||
if optimal:
|
||||
path = os.path.join(self.args.save_dir, 'results', 'model')
|
||||
else:
|
||||
path = os.path.join(self.model_save_dir, f'{self.epoch_cnt:04d}')
|
||||
|
||||
os.makedirs(path, exist_ok=True)
|
||||
nemf = self.field.module if isinstance(self.field, nn.DataParallel) else self.field
|
||||
torch.save(nemf.state_dict(), os.path.join(path, 'nemf.pth'))
|
||||
torch.save(self.optimizer.state_dict(), os.path.join(path, 'optimizer.pth'))
|
||||
if self.args.scheduler.name:
|
||||
torch.save(self.schedulers[0].state_dict(), os.path.join(path, 'scheduler.pth'))
|
||||
self.loss_recorder.save(path)
|
||||
|
||||
print(f'Save at {path} succeeded')
|
||||
|
||||
def load(self, epoch=None, optimal=False):
|
||||
if optimal:
|
||||
path = os.path.join(self.args.save_dir, 'results', 'model')
|
||||
else:
|
||||
if epoch is None:
|
||||
all = [int(q) for q in os.listdir(self.model_save_dir)]
|
||||
if len(all) == 0:
|
||||
raise RuntimeError(f'Empty loading path {self.model_save_dir}')
|
||||
epoch = sorted(all)[-1]
|
||||
path = os.path.join(self.model_save_dir, f'{epoch:04d}')
|
||||
|
||||
print(f'Loading from {path}')
|
||||
nemf = self.field.module if isinstance(self.field, nn.DataParallel) else self.field
|
||||
nemf.load_state_dict(torch.load(os.path.join(path, 'nemf.pth'), map_location=self.device))
|
||||
if self.is_train:
|
||||
self.optimizer.load_state_dict(torch.load(os.path.join(path, 'optimizer.pth')))
|
||||
self.loss_recorder.load(path)
|
||||
if self.args.scheduler.name:
|
||||
self.schedulers[0].load_state_dict(torch.load(os.path.join(path, 'scheduler.pth')))
|
||||
self.epoch_cnt = epoch if not optimal else 0
|
||||
|
||||
print('Load succeeded')
|
||||
|
||||
def super_sampling(self, step=1):
|
||||
with torch.no_grad():
|
||||
self.forward(step)
|
||||
self.fps = self.args.data.fps / step
|
||||
self.compute_test_result()
|
||||
|
||||
def compute_test_result(self):
|
||||
os.makedirs(self.viz_dir, exist_ok=True)
|
||||
|
||||
rotmat = self.recon_data['rotmat'][0] # (T, J, 3, 3)
|
||||
rotmat_gt = self.input_data['rotmat'][0] # (T, J, 3, 3)
|
||||
local_rotmat = self.fk.global_to_local(rotmat) # (T, J, 3, 3)
|
||||
local_rotmat_gt = self.fk.global_to_local(rotmat_gt) # (T, J, 3, 3)
|
||||
|
||||
if self.args.data.root_transform:
|
||||
root_orient = rotation_6d_to_matrix(self.recon_data['root_orient'][0]) # (T, 3, 3)
|
||||
root_orient_gt = rotation_6d_to_matrix(self.input_data['root_orient'][0]) # (T, 3, 3)
|
||||
local_rotmat[:, 0] = root_orient
|
||||
local_rotmat_gt[:, 0] = root_orient_gt
|
||||
rotation = matrix_to_quaternion(local_rotmat) # (T, J, 4)
|
||||
rotation_gt = matrix_to_quaternion(local_rotmat_gt) # (T, J, 4)
|
||||
|
||||
origin = self.input_data['trans'][0, 0].clone() # [3]
|
||||
position = self.recon_data['trans'][0].clone() # (T, 3)
|
||||
position[:, self.v_axis] = position[:, self.v_axis] + origin[self.v_axis].unsqueeze(0)
|
||||
position_gt = self.input_data['trans'][0] # (T, 3)
|
||||
|
||||
if self.args.bvh_viz:
|
||||
if self.args.amass_data:
|
||||
rotation = align_joints(rotation)
|
||||
rotation_gt = align_joints(rotation_gt)
|
||||
|
||||
position = position.unsqueeze(1) # (T, 1, 3)
|
||||
position_gt = position_gt.unsqueeze(1) # (T, 1, 3)
|
||||
|
||||
# export data to BVH files
|
||||
if self.args.amass_data:
|
||||
offsets = self.args.smpl.offsets[self.args.data.gender]
|
||||
parents = self.args.smpl.parents
|
||||
names = self.args.smpl.joint_names
|
||||
else:
|
||||
offsets = self.args.animal.offsets
|
||||
parents = self.args.animal.parents
|
||||
names = self.args.animal.joint_names
|
||||
|
||||
anim = Animation(Quaternions(c2c(rotation)), c2c(position), None, offsets=offsets, parents=parents)
|
||||
BVH.save(os.path.join(self.viz_dir, f'test_{int(self.fps)}fps.bvh'), anim, names=names, frametime=1 / self.fps)
|
||||
|
||||
anim.rotations = Quaternions(c2c(rotation_gt))
|
||||
anim.positions = c2c(position_gt)
|
||||
BVH.save(os.path.join(self.viz_dir, 'test_gt.bvh'), anim, names=names, frametime=1 / self.args.data.fps)
|
||||
else:
|
||||
poses = c2c(quaternion_to_axis_angle(rotation)) # (T, J, 3)
|
||||
poses_gt = c2c(quaternion_to_axis_angle(rotation_gt)) # (T, J, 3)
|
||||
|
||||
poses = poses.reshape((poses.shape[0], -1)) # (T, J x 3)
|
||||
poses = np.pad(poses, [(0, 0), (0, 93)], mode='constant')
|
||||
poses_gt = poses_gt.reshape((poses_gt.shape[0], -1)) # (T, J x 3)
|
||||
poses_gt = np.pad(poses_gt, [(0, 0), (0, 93)], mode='constant')
|
||||
|
||||
np.savez(os.path.join(self.viz_dir, f'test_{int(self.fps)}fps.npz'),
|
||||
poses=poses, trans=c2c(position), betas=np.zeros(10), gender=self.args.data.gender, mocap_framerate=self.fps)
|
||||
np.savez(os.path.join(self.viz_dir, 'test_gt.npz'),
|
||||
poses=poses_gt, trans=c2c(position_gt), betas=np.zeros(10), gender=self.args.data.gender, mocap_framerate=self.args.data.fps)
|
||||
+152
@@ -0,0 +1,152 @@
|
||||
"""Based on Daniel Holden code from:
|
||||
A Deep Learning Framework for Character Motion Synthesis and Editing
|
||||
(http://www.ipab.inf.ed.ac.uk/cgvu/motionsynthesis.pdf)
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from rotations import euler_angles_to_matrix, quaternion_to_matrix, rotation_6d_to_matrix
|
||||
|
||||
|
||||
class ForwardKinematicsLayer(nn.Module):
|
||||
""" Forward Kinematics Layer Class """
|
||||
|
||||
def __init__(self, args=None, parents=None, positions=None, device=None):
|
||||
super().__init__()
|
||||
self.b_idxs = None
|
||||
if device is None:
|
||||
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||||
else:
|
||||
self.device = device
|
||||
if parents is None and positions is None:
|
||||
# Load SMPL skeleton (their joint order is different from the one we use for bvh export)
|
||||
# smpl_fname = os.path.join(args.smpl.smpl_body_model, args.data.gender, 'model.npz')
|
||||
smpl_fname = os.path.join(args.smpl.smpl_body_model, "MANO_RIGHT.pkl")
|
||||
if not os.path.exists(smpl_fname):
|
||||
raise FileNotFoundError(f'SMPL file {os.path.abspath(smpl_fname)} does not exist. Please check the path., pwd: {os.getcwd()}')
|
||||
smpl_data = np.load(smpl_fname, encoding='latin1', allow_pickle=True)
|
||||
self.parents = torch.from_numpy(smpl_data['kintree_table'][0].astype(np.int32)).to(self.device)
|
||||
self.parents = self.parents.long()
|
||||
self.positions = torch.from_numpy(smpl_data['J'].astype(np.float32)).to(self.device)
|
||||
self.positions[1:] -= self.positions[self.parents[1:]]
|
||||
else:
|
||||
self.parents = torch.from_numpy(parents).to(self.device)
|
||||
self.parents = self.parents.long()
|
||||
self.positions = torch.from_numpy(positions).to(self.device)
|
||||
self.positions = self.positions.float()
|
||||
self.positions[0] = 0
|
||||
|
||||
def rotate(self, t0s, t1s):
|
||||
return torch.matmul(t0s, t1s)
|
||||
|
||||
def identity_rotation(self, rotations):
|
||||
diagonal = torch.diag(torch.tensor([1.0, 1.0, 1.0, 1.0])).to(self.device)
|
||||
diagonal = torch.reshape(
|
||||
diagonal, torch.Size([1] * len(rotations.shape[:2]) + [4, 4]))
|
||||
ts = diagonal.repeat(rotations.shape[:2] + torch.Size([1, 1]))
|
||||
return ts
|
||||
|
||||
def make_fast_rotation_matrices(self, positions, rotations):
|
||||
if len(rotations.shape) == 4 and rotations.shape[-2:] == torch.Size([3, 3]):
|
||||
rot_matrices = rotations
|
||||
elif rotations.shape[-1] == 3:
|
||||
rot_matrices = euler_angles_to_matrix(rotations, convention='XYZ')
|
||||
elif rotations.shape[-1] == 4:
|
||||
rot_matrices = quaternion_to_matrix(rotations)
|
||||
elif rotations.shape[-1] == 6:
|
||||
rot_matrices = rotation_6d_to_matrix(rotations)
|
||||
else:
|
||||
raise NotImplementedError(f'Unimplemented rotation representation in FK layer, shape of {rotations.shape}')
|
||||
|
||||
rot_matrices = torch.cat([rot_matrices, positions[..., None]], dim=-1)
|
||||
zeros = torch.zeros(rot_matrices.shape[:-2] + torch.Size([1, 3])).to(self.device)
|
||||
ones = torch.ones(rot_matrices.shape[:-2] + torch.Size([1, 1])).to(self.device)
|
||||
zerosones = torch.cat([zeros, ones], dim=-1)
|
||||
rot_matrices = torch.cat([rot_matrices, zerosones], dim=-2)
|
||||
return rot_matrices
|
||||
|
||||
def rotate_global(self, parents, positions, rotations):
|
||||
locals = self.make_fast_rotation_matrices(positions, rotations)
|
||||
globals = self.identity_rotation(rotations)
|
||||
|
||||
globals = torch.cat([locals[:, 0:1], globals[:, 1:]], dim=1)
|
||||
b_size = positions.shape[0]
|
||||
if self.b_idxs is None:
|
||||
self.b_idxs = torch.LongTensor(np.arange(b_size)).to(self.device)
|
||||
elif self.b_idxs.shape[-1] != b_size:
|
||||
self.b_idxs = torch.LongTensor(np.arange(b_size)).to(self.device)
|
||||
|
||||
for i in range(1, positions.shape[1]):
|
||||
globals[:, i] = self.rotate(
|
||||
globals[self.b_idxs, parents[i]], locals[:, i])
|
||||
|
||||
return globals
|
||||
|
||||
def get_tpose_joints(self, offsets, parents):
|
||||
num_joints = len(parents)
|
||||
joints = [offsets[:, 0]]
|
||||
for j in range(1, len(parents)):
|
||||
joints.append(joints[parents[j]] + offsets[:, j])
|
||||
|
||||
return torch.stack(joints, dim=1)
|
||||
|
||||
def canonical_to_local(self, canonical_xform, global_orient=None):
|
||||
"""
|
||||
Args:
|
||||
canonical_xform: (B, J, 3, 3)
|
||||
global_orient: (B, 3, 3)
|
||||
|
||||
Returns:
|
||||
local_xform: (B, J, 3, 3)
|
||||
"""
|
||||
local_xform = torch.zeros_like(canonical_xform)
|
||||
|
||||
if global_orient is None:
|
||||
global_xform = canonical_xform
|
||||
else:
|
||||
global_xform = torch.matmul(global_orient.unsqueeze(1), canonical_xform)
|
||||
for i in range(global_xform.shape[1]):
|
||||
if i == 0:
|
||||
local_xform[:, i] = global_xform[:, i]
|
||||
else:
|
||||
local_xform[:, i] = torch.bmm(torch.linalg.inv(global_xform[:, self.parents[i]]), global_xform[:, i])
|
||||
|
||||
return local_xform
|
||||
|
||||
def global_to_local(self, global_xform):
|
||||
"""
|
||||
Args:
|
||||
global_xform: (B, J, 3, 3)
|
||||
|
||||
Returns:
|
||||
local_xform: (B, J, 3, 3)
|
||||
"""
|
||||
local_xform = torch.zeros_like(global_xform)
|
||||
|
||||
for i in range(global_xform.shape[1]):
|
||||
if i == 0:
|
||||
local_xform[:, i] = global_xform[:, i]
|
||||
else:
|
||||
local_xform[:, i] = torch.bmm(torch.linalg.inv(global_xform[:, self.parents[i]]), global_xform[:, i])
|
||||
|
||||
return local_xform
|
||||
|
||||
def forward(self, rotations, positions=None):
|
||||
"""
|
||||
Args:
|
||||
rotations (B, J, D)
|
||||
|
||||
Returns:
|
||||
The global position of each joint after FK (B, J, 3)
|
||||
"""
|
||||
# Get the full transform with rotations for skinning
|
||||
b_size = rotations.shape[0]
|
||||
if positions is None:
|
||||
positions = self.positions.repeat(b_size, 1, 1)
|
||||
transforms = self.rotate_global(self.parents, positions, rotations)
|
||||
coordinates = transforms[:, :, :3, 3] / transforms[:, :, 3:, 3]
|
||||
|
||||
return coordinates, transforms
|
||||
@@ -0,0 +1,525 @@
|
||||
import os
|
||||
import yaml
|
||||
|
||||
import holden.BVH as BVH
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from holden.Animation import Animation
|
||||
from holden.Quaternions import Quaternions
|
||||
from human_body_prior.tools.omni_tools import copy2cpu as c2c
|
||||
from rotations import matrix_to_quaternion, matrix_to_rotation_6d, quaternion_to_axis_angle, rotation_6d_to_matrix
|
||||
from utils import align_joints, compute_trajectory, estimate_angular_velocity, estimate_linear_velocity, normalize
|
||||
|
||||
from .base_model import BaseModel
|
||||
from .fk import ForwardKinematicsLayer
|
||||
from .global_motion import GlobalMotionPredictor
|
||||
from .losses import GeodesicLoss
|
||||
from .neural_motion import NeuralMotionField
|
||||
from .prior import GlobalEncoder, LocalEncoder
|
||||
from .skeleton import build_edge_topology
|
||||
|
||||
|
||||
|
||||
class Struct:
|
||||
def __init__(self, **entries):
|
||||
self.__dict__.update(entries)
|
||||
|
||||
class Arguments:
|
||||
def __init__(self, config_path, filename='default.yaml'):
|
||||
with open(os.path.join(config_path, 'mano.yaml'), 'r') as f:
|
||||
smpl = yaml.safe_load(f)
|
||||
|
||||
self.smpl = Struct(**smpl)
|
||||
self.smpl.offsets['right'] = np.array(self.smpl.offsets['right'])
|
||||
self.smpl.offsets['left'] = np.array(self.smpl.offsets['left'])
|
||||
self.smpl.parents = np.array(self.smpl.parents).astype(np.int32)
|
||||
self.smpl.joint_num = 16 # len(self.smpl.joints_to_use)
|
||||
self.smpl.joints_to_use = np.array(self.smpl.joints_to_use)
|
||||
self.smpl.joints_to_use = np.arange(0, 63).reshape((-1, 3))[self.smpl.joints_to_use].reshape(-1)
|
||||
|
||||
self.filename = os.path.splitext(filename)[0]
|
||||
with open(os.path.join(config_path, filename), 'r') as f:
|
||||
config = yaml.safe_load(f)
|
||||
|
||||
for key, value in config.items():
|
||||
if isinstance(value, dict):
|
||||
setattr(self, key, Struct(**value))
|
||||
else:
|
||||
setattr(self, key, value)
|
||||
|
||||
self.json = config
|
||||
|
||||
|
||||
class Architecture(BaseModel):
|
||||
def __init__(self, args, ngpu, batch_num=0):
|
||||
super(Architecture, self).__init__(args)
|
||||
|
||||
self.batch_num = batch_num
|
||||
self.args = args
|
||||
#from pprint import pprint
|
||||
#pprint(vars(self.args))
|
||||
|
||||
self.fk = ForwardKinematicsLayer(args)
|
||||
|
||||
smpl_fname = os.path.join(args.smpl.smpl_body_model, "MANO_RIGHT.pkl")
|
||||
smpl_data = np.load(smpl_fname, allow_pickle=True, encoding='latin1')
|
||||
parents = smpl_data['kintree_table'][0].astype(np.int32)
|
||||
edges = build_edge_topology(parents)
|
||||
|
||||
self.local_encoder = LocalEncoder(args.local_prior, edges).to(self.device)
|
||||
self.global_encoder = GlobalEncoder(args.global_prior).to(self.device)
|
||||
self.field = NeuralMotionField(args.nemf).to(self.device)
|
||||
|
||||
if args.multi_gpu is True:
|
||||
self.local_encoder = nn.DataParallel(self.local_encoder, device_ids=range(ngpu))
|
||||
self.global_encoder = nn.DataParallel(self.global_encoder, device_ids=range(ngpu))
|
||||
self.field = nn.DataParallel(self.field, device_ids=range(ngpu))
|
||||
|
||||
self.models = [self.local_encoder, self.field]
|
||||
if args.data.root_transform:
|
||||
self.models.append(self.global_encoder)
|
||||
|
||||
ordermap = {
|
||||
'x': 0,
|
||||
'y': 1,
|
||||
'z': 2,
|
||||
}
|
||||
self.up_axis = ordermap[self.args.data.up]
|
||||
self.v_axis = [x for x in ordermap.values() if x != self.up_axis]
|
||||
|
||||
print(self.up_axis)
|
||||
print(self.v_axis)
|
||||
|
||||
self.input_data = dict()
|
||||
self.recon_data = dict()
|
||||
print(os.getcwd(), '...', os.path.join(args.dataset_dir, f'mean-{args.data.gender}-{args.data.clip_length}-{args.data.fps}fps.pt'))
|
||||
self.mean = torch.load(os.path.join(args.dataset_dir, f'mean-{args.data.gender}-{args.data.clip_length}-{args.data.fps}fps.pt'), map_location=self.device)
|
||||
print(f'mean: {list(self.mean.keys())}')
|
||||
self.std = torch.load(os.path.join(args.dataset_dir, f'std-{args.data.gender}-{args.data.clip_length}-{args.data.fps}fps.pt'), map_location=self.device)
|
||||
print(f'std: {list(self.std.keys())}')
|
||||
|
||||
#for k, v in self.mean.items():
|
||||
# print(k, v.shape)
|
||||
|
||||
if args.pretrained_gmp:
|
||||
assert args.output_trans is False, 'already output global translation in this model'
|
||||
gmp_args = Arguments('./configs', filename=args.pretrained_gmp)
|
||||
self.gmp = GlobalMotionPredictor(gmp_args, ngpu)
|
||||
self.gmp.load(optimal=True)
|
||||
self.gmp.eval()
|
||||
|
||||
self.criterion_geo = GeodesicLoss().to(self.device)
|
||||
if self.is_train:
|
||||
parameters = list(self.local_encoder.parameters()) + list(self.field.parameters())
|
||||
if args.data.root_transform:
|
||||
parameters += list(self.global_encoder.parameters())
|
||||
|
||||
if args.adam_optimizer:
|
||||
self.optimizer = torch.optim.Adam(parameters, args.learning_rate, weight_decay=args.weight_decay)
|
||||
else:
|
||||
self.optimizer = torch.optim.Adamax(parameters, args.learning_rate, weight_decay=args.weight_decay)
|
||||
self.optimizers = [self.optimizer]
|
||||
|
||||
self.criterion_rec = nn.L1Loss() if args.l1_loss else nn.MSELoss()
|
||||
self.criterion_bce = nn.BCEWithLogitsLoss()
|
||||
|
||||
self.iteration = 0
|
||||
else:
|
||||
self.fps = args.data.fps
|
||||
self.test_index = 0
|
||||
if args.bvh_viz:
|
||||
self.viz_dir = os.path.join(args.save_dir, 'results', 'bvh')
|
||||
else:
|
||||
self.viz_dir = os.path.join(args.save_dir, 'results', 'smpl')
|
||||
|
||||
def set_input(self, input):
|
||||
self.input_data = {k: v.float().to(self.device) for k, v in input.items() if k in [
|
||||
'pos', 'velocity', 'global_xform', 'angular', 'root_orient', 'root_vel', 'trans', 'contacts'
|
||||
]}
|
||||
self.input_data['rotmat'] = rotation_6d_to_matrix(self.input_data['global_xform'])
|
||||
|
||||
def reparameterize(self, mu, logvar):
|
||||
eps = torch.randn_like(mu)
|
||||
std = torch.exp(0.5 * logvar)
|
||||
return mu + eps * std
|
||||
|
||||
def encode_local(self):
|
||||
b_size, t_length = self.input_data['pos'].shape[:2]
|
||||
|
||||
if 'pos' in self.args.data.normalize:
|
||||
pos = normalize(self.input_data['pos'], mean=self.mean['pos'], std=self.std['pos'])
|
||||
else:
|
||||
pos = self.input_data['pos']
|
||||
if 'velocity' in self.args.data.normalize:
|
||||
velocity = normalize(self.input_data['velocity'], mean=self.mean['velocity'], std=self.std['velocity'])
|
||||
else:
|
||||
velocity = self.input_data['velocity']
|
||||
if 'global_xform' in self.args.data.normalize:
|
||||
global_xform = normalize(self.input_data['global_xform'], mean=self.mean['global_xform'], std=self.std['global_xform'])
|
||||
else:
|
||||
global_xform = self.input_data['global_xform']
|
||||
if 'angular' in self.args.data.normalize:
|
||||
angular = normalize(self.input_data['angular'], mean=self.mean['angular'], std=self.std['angular'])
|
||||
else:
|
||||
angular = self.input_data['angular']
|
||||
|
||||
x = torch.cat((pos, velocity, global_xform, angular), dim=-1)
|
||||
x = x.view(b_size, t_length, -1) # (B, T, J x D)
|
||||
x = x.permute(0, 2, 1) # (B, J x D, T)
|
||||
|
||||
mu, logvar = self.local_encoder(x) # (B, D)
|
||||
if self.args.lambda_kl != 0:
|
||||
z = self.reparameterize(mu, logvar)
|
||||
else:
|
||||
z = mu
|
||||
|
||||
return z, mu, logvar
|
||||
|
||||
def encode_global(self):
|
||||
if 'root_orient' in self.args.data.normalize:
|
||||
root_orient = normalize(self.input_data['root_orient'], mean=self.mean['root_orient'], std=self.std['root_orient'])
|
||||
else:
|
||||
root_orient = self.input_data['root_orient']
|
||||
if self.args.global_prior.in_channels == 9:
|
||||
if 'root_vel' in self.args.data.normalize:
|
||||
root_vel = normalize(self.input_data['root_vel'], mean=self.mean['root_vel'], std=self.std['root_vel'])
|
||||
else:
|
||||
root_vel = self.input_data['root_vel']
|
||||
x = torch.cat((root_orient, root_vel), dim=-1) # (B, T, D)
|
||||
elif self.args.global_prior.in_channels == 6:
|
||||
x = root_orient
|
||||
else:
|
||||
raise RuntimeError('unsupported input dimensions for global prior')
|
||||
x = x.permute(0, 2, 1) # (B, D, T)
|
||||
|
||||
mu, logvar = self.global_encoder(x) # (B, D)
|
||||
if self.args.lambda_kl != 0:
|
||||
z = self.reparameterize(mu, logvar)
|
||||
else:
|
||||
z = mu
|
||||
|
||||
return z, mu, logvar
|
||||
|
||||
def forward(self, step=1):
|
||||
z_l, mu_l, logvar_l = self.encode_local() # (B, D)
|
||||
if self.args.data.root_transform:
|
||||
z_g, mu_g, logvar_g = self.encode_global() # (B, D)
|
||||
else:
|
||||
z_g = None
|
||||
|
||||
self.recon_data.clear()
|
||||
self.recon_data = self.decode(z_l, z_g, length=self.args.data.clip_length, step=step)
|
||||
|
||||
self.recon_data['mu_l'] = mu_l
|
||||
self.recon_data['logvar_l'] = logvar_l
|
||||
if self.args.data.root_transform:
|
||||
self.recon_data['mu_g'] = mu_g
|
||||
self.recon_data['logvar_g'] = logvar_g
|
||||
|
||||
def decode(self, z_l, z_g, length, step=1):
|
||||
b_size = z_l.shape[0]
|
||||
n_joints = self.args.smpl.joint_num
|
||||
n_joints = 16
|
||||
|
||||
# generate f(t; z)
|
||||
t = torch.arange(start=0, end=length, step=step).unsqueeze(0) # (1, T)
|
||||
t = t / self.args.data.clip_length * 2 - 1 # map t to [-1, 1]
|
||||
t = t.expand(b_size, -1).unsqueeze(-1).to(self.device) # (B, T, 1)
|
||||
local_motion, global_motion = self.field(t, z_l, z_g) # (B, T, N)
|
||||
|
||||
rot6d_recon = local_motion[:, :, :n_joints * 6].contiguous() # (B, T, J x 6)
|
||||
rot6d_recon = rot6d_recon.view(b_size, -1, n_joints, 6) # (B, T, J, 6)
|
||||
rotmat_recon = rotation_6d_to_matrix(rot6d_recon) # (B, T, J, 3, 3)
|
||||
|
||||
if self.args.data.root_transform:
|
||||
identity = torch.zeros(b_size, t.shape[1], 3, 3).to(self.device)
|
||||
identity[:, :, 0, 0] = 1
|
||||
identity[:, :, 1, 1] = 1
|
||||
identity[:, :, 2, 2] = 1
|
||||
rotmat_recon[:, :, 0] = identity
|
||||
# rotmat_recon[:, :, -2] = rotmat_recon[:, :, -4]
|
||||
# rotmat_recon[:, :, -1] = rotmat_recon[:, :, -3]
|
||||
local_rotmat = self.fk.global_to_local(rotmat_recon.view(-1, n_joints, 3, 3)) # (B x T, J, 3. 3)
|
||||
pos_recon, _ = self.fk(local_rotmat) # (B x T, J, 3)
|
||||
pos_recon = pos_recon.contiguous().view(b_size, -1, n_joints, 3) # (B, T, J, 3)
|
||||
|
||||
output = dict()
|
||||
output['rotmat'] = rotmat_recon
|
||||
output['pos'] = pos_recon
|
||||
|
||||
if self.args.data.root_transform:
|
||||
root_orient = global_motion[:, :, :6] # (B, T, 6)
|
||||
output['root_orient'] = root_orient
|
||||
else:
|
||||
output['root_orient'] = rot6d_recon[:, :, 0] # (B, T, 6)
|
||||
|
||||
dt = 1.0 / self.args.data.fps * step
|
||||
if self.args.output_trans:
|
||||
root_vel = global_motion[:, :, 6:9] # (B, T, 3)
|
||||
root_height = global_motion[:, :, 9] # (B, T)
|
||||
|
||||
output['root_vel'] = root_vel
|
||||
output['root_height'] = root_height
|
||||
|
||||
contacts = global_motion[:, :, 10:] # (B, T, 8)
|
||||
output['contacts'] = contacts
|
||||
|
||||
origin = torch.zeros(b_size, 3).to(self.device)
|
||||
trans = compute_trajectory(root_vel, root_height, origin, dt, up_axis=self.args.data.up)
|
||||
output['trans'] = trans
|
||||
elif self.args.pretrained_gmp:
|
||||
local_rotmat = local_rotmat.view(b_size, -1, n_joints, 3, 3) # (B, T, J, 3, 3)
|
||||
if self.args.data.root_transform:
|
||||
root_orient_rotmat = rotation_6d_to_matrix(output['root_orient']) # (B, T, 3, 3)
|
||||
local_rotmat[:, :, 0] = root_orient_rotmat
|
||||
pos = torch.matmul(root_orient_rotmat.unsqueeze(2), pos_recon.unsqueeze(-1)).squeeze(-1)
|
||||
else:
|
||||
pos = pos_recon
|
||||
|
||||
gmp_data = dict()
|
||||
gmp_data['rot6d'] = matrix_to_rotation_6d(local_rotmat) # (B, T, J, 6)
|
||||
gmp_data['angular'] = estimate_angular_velocity(local_rotmat.clone(), dt=dt) # (B, T, J, 3)
|
||||
|
||||
gmp_data['pos'] = pos
|
||||
gmp_data['velocity'] = estimate_linear_velocity(pos, dt=dt)
|
||||
|
||||
pred_data = self.gmp.predict(gmp_data, dt=dt)
|
||||
|
||||
output |= pred_data # merge two dicts
|
||||
|
||||
return output
|
||||
|
||||
def kl_scheduler(self):
|
||||
"""
|
||||
Cyclical Annealing Schedule
|
||||
"""
|
||||
if (self.epoch_cnt % self.args.annealing_cycles == 0) and (self.iteration % self.batch_num == 0):
|
||||
print('KL annealing restart')
|
||||
return 0.01
|
||||
|
||||
return min(1, 0.01 + self.iteration / (self.args.annealing_warmup * self.batch_num))
|
||||
|
||||
def backward(self, validation=False):
|
||||
loss = 0
|
||||
|
||||
root_orient = rotation_6d_to_matrix(self.recon_data['root_orient'])
|
||||
root_orient_gt = rotation_6d_to_matrix(self.input_data['root_orient'])
|
||||
|
||||
if self.args.geodesic_loss:
|
||||
orient_recon_loss = self.criterion_geo(root_orient.view(-1, 3, 3), root_orient_gt.view(-1, 3, 3))
|
||||
else:
|
||||
orient_recon_loss = self.criterion_rec(root_orient, root_orient_gt)
|
||||
self.loss_recorder.add_scalar('orient_recon_loss', orient_recon_loss, validation=validation)
|
||||
loss += self.args.lambda_orient * orient_recon_loss
|
||||
|
||||
if self.args.geodesic_loss:
|
||||
rotmat_recon_loss = self.criterion_geo(self.recon_data['rotmat'].view(-1, 3, 3), self.input_data['rotmat'].view(-1, 3, 3))
|
||||
else:
|
||||
rotmat_recon_loss = self.criterion_rec(self.recon_data['rotmat'], self.input_data['rotmat'])
|
||||
self.loss_recorder.add_scalar('rotmat_recon_loss', rotmat_recon_loss, validation=validation)
|
||||
loss += self.args.lambda_rotmat * rotmat_recon_loss
|
||||
|
||||
pos_recon_loss = self.criterion_rec(self.recon_data['pos'], self.input_data['pos'])
|
||||
self.loss_recorder.add_scalar('pos_recon_loss', pos_recon_loss, validation=validation)
|
||||
loss += self.args.lambda_pos * pos_recon_loss
|
||||
|
||||
if self.args.output_trans:
|
||||
v_loss = self.criterion_rec(self.recon_data['root_vel'], self.input_data['root_vel'])
|
||||
up_loss = self.criterion_rec(self.recon_data['root_height'], self.input_data['trans'][..., 'xyz'.index(self.args.data.up)])
|
||||
self.loss_recorder.add_scalar('v_loss', v_loss, validation=validation)
|
||||
self.loss_recorder.add_scalar('up_loss', up_loss, validation=validation)
|
||||
loss += self.args.lambda_v * v_loss + self.args.lambda_up * up_loss
|
||||
|
||||
origin = self.input_data['trans'][:, 0]
|
||||
trans = self.recon_data['trans']
|
||||
trans[..., self.v_axis] = trans[..., self.v_axis] + origin[..., self.v_axis].unsqueeze(1)
|
||||
trans_loss = self.criterion_rec(trans, self.input_data['trans'])
|
||||
self.loss_recorder.add_scalar('trans_loss', trans_loss, validation=validation)
|
||||
loss += self.args.lambda_trans * trans_loss
|
||||
|
||||
contacts_loss = self.criterion_bce(self.recon_data['contacts'], self.input_data['contacts'])
|
||||
self.loss_recorder.add_scalar('contacts_loss', contacts_loss, validation=validation)
|
||||
loss += self.args.lambda_contacts * contacts_loss
|
||||
|
||||
if self.args.lambda_kl != 0:
|
||||
mu_l = self.recon_data['mu_l']
|
||||
logvar_l = self.recon_data['logvar_l']
|
||||
local_kl_loss = torch.mean(-0.5 * torch.sum(1 + logvar_l - mu_l.pow(2) - logvar_l.exp(), dim=-1))
|
||||
self.loss_recorder.add_scalar('local_kl_loss', local_kl_loss, validation=validation)
|
||||
|
||||
mu_g = self.recon_data['mu_g']
|
||||
logvar_g = self.recon_data['logvar_g']
|
||||
global_kl_loss = torch.mean(-0.5 * torch.sum(1 + logvar_g - mu_g.pow(2) - logvar_g.exp(), dim=-1))
|
||||
self.loss_recorder.add_scalar('global_kl_loss', global_kl_loss, validation=validation)
|
||||
if not validation:
|
||||
loss += self.args.lambda_kl * self.kl_scheduler() * (local_kl_loss + global_kl_loss)
|
||||
|
||||
self.loss_recorder.add_scalar('total_loss', loss, validation=validation)
|
||||
|
||||
if not validation:
|
||||
loss.backward()
|
||||
|
||||
def optimize_parameters(self):
|
||||
self.optimizer.zero_grad()
|
||||
self.forward()
|
||||
self.backward(validation=False)
|
||||
self.optimizer.step()
|
||||
|
||||
self.iteration += 1
|
||||
|
||||
def report_errors(self):
|
||||
root_orient = rotation_6d_to_matrix(self.recon_data['root_orient']) # (B, T, 3, 3)
|
||||
root_orient_gt = rotation_6d_to_matrix(self.input_data['root_orient']) # (B, T, 3, 3)
|
||||
orientation_error = self.criterion_geo(root_orient.view(-1, 3, 3), root_orient_gt.view(-1, 3, 3))
|
||||
|
||||
local_rotmat = self.fk.global_to_local(self.recon_data['rotmat'].view(-1, self.args.smpl.joint_num, 3, 3)) # (B x T, J, 3, 3)
|
||||
local_rotmat_gt = self.fk.global_to_local(self.input_data['rotmat'].view(-1, self.args.smpl.joint_num, 3, 3)) # (B x T, J, 3, 3)
|
||||
rotation_error = self.criterion_geo(local_rotmat[:, 1:].reshape(-1, 3, 3), local_rotmat_gt[:, 1:].reshape(-1, 3, 3))
|
||||
|
||||
pos = self.recon_data['pos'] # (B, T, J, 3)
|
||||
pos_gt = self.input_data['pos'] # (B, T, J, 3)
|
||||
position_error = torch.linalg.norm((pos - pos_gt), dim=-1).mean()
|
||||
|
||||
if 'trans' in self.recon_data.keys():
|
||||
origin = self.input_data['trans'][:, 0]
|
||||
trans = self.recon_data['trans']
|
||||
trans[..., self.v_axis] = trans[..., self.v_axis] + origin[..., self.v_axis].unsqueeze(1)
|
||||
trans_gt = self.input_data['trans'] # (B, T, 3)
|
||||
translation_error = torch.linalg.norm((trans - trans_gt), dim=-1).mean()
|
||||
else:
|
||||
translation_error = torch.zeros(1).mean()
|
||||
|
||||
return {
|
||||
'rotation': c2c(rotation_error),
|
||||
'position': c2c(position_error),
|
||||
'orientation': c2c(orientation_error),
|
||||
'translation': c2c(translation_error)
|
||||
}
|
||||
|
||||
def verbose(self):
|
||||
res = {}
|
||||
for loss in self.loss_recorder.losses.values():
|
||||
res[loss.name] = {'train': loss.current()[0], 'val': loss.current()[1]}
|
||||
|
||||
return res
|
||||
|
||||
def validate(self):
|
||||
with torch.no_grad():
|
||||
self.forward()
|
||||
self.backward(validation=True)
|
||||
|
||||
def save(self, optimal=False):
|
||||
if optimal:
|
||||
path = os.path.join(self.args.save_dir, 'results', 'model')
|
||||
else:
|
||||
path = os.path.join(self.model_save_dir, f'{self.epoch_cnt:04d}')
|
||||
|
||||
os.makedirs(path, exist_ok=True)
|
||||
local_encoder = self.local_encoder.module if isinstance(self.local_encoder, nn.DataParallel) else self.local_encoder
|
||||
nemf = self.field.module if isinstance(self.field, nn.DataParallel) else self.field
|
||||
torch.save(local_encoder.state_dict(), os.path.join(path, 'local_encoder.pth'))
|
||||
torch.save(nemf.state_dict(), os.path.join(path, 'nemf.pth'))
|
||||
if self.args.data.root_transform:
|
||||
global_encoder = self.global_encoder.module if isinstance(self.global_encoder, nn.DataParallel) else self.global_encoder
|
||||
torch.save(global_encoder.state_dict(), os.path.join(path, 'global_encoder.pth'))
|
||||
torch.save(self.optimizer.state_dict(), os.path.join(path, 'optimizer.pth'))
|
||||
if self.args.scheduler.name:
|
||||
torch.save(self.schedulers[0].state_dict(), os.path.join(path, 'scheduler.pth'))
|
||||
self.loss_recorder.save(path)
|
||||
|
||||
print(f'Save at {path} succeeded')
|
||||
|
||||
def load(self, epoch=None, optimal=False):
|
||||
if optimal:
|
||||
path = os.path.join(self.args.save_dir, 'results', 'model')
|
||||
else:
|
||||
if epoch is None:
|
||||
all = [int(q) for q in os.listdir(self.model_save_dir)]
|
||||
if len(all) == 0:
|
||||
raise RuntimeError(f'Empty loading path {self.model_save_dir}')
|
||||
epoch = sorted(all)[-1]
|
||||
path = os.path.join(self.model_save_dir, f'{epoch:04d}')
|
||||
|
||||
print(f'Loading from {path}')
|
||||
local_encoder = self.local_encoder.module if isinstance(self.local_encoder, nn.DataParallel) else self.local_encoder
|
||||
nemf = self.field.module if isinstance(self.field, nn.DataParallel) else self.field
|
||||
local_encoder.load_state_dict(torch.load(os.path.join(path, 'local_encoder.pth'), map_location=self.device))
|
||||
nemf.load_state_dict(torch.load(os.path.join(path, 'nemf.pth'), map_location=self.device))
|
||||
if self.args.data.root_transform:
|
||||
global_encoder = self.global_encoder.module if isinstance(self.global_encoder, nn.DataParallel) else self.global_encoder
|
||||
global_encoder.load_state_dict(torch.load(os.path.join(path, 'global_encoder.pth'), map_location=self.device))
|
||||
if self.is_train:
|
||||
self.optimizer.load_state_dict(torch.load(os.path.join(path, 'optimizer.pth')))
|
||||
self.loss_recorder.load(path)
|
||||
self.iteration = len(self.loss_recorder.losses['total_loss'].loss_step)
|
||||
print(f'trained iterations: {self.iteration}')
|
||||
if self.args.scheduler.name:
|
||||
self.schedulers[0].load_state_dict(torch.load(os.path.join(path, 'scheduler.pth')))
|
||||
self.epoch_cnt = epoch if not optimal else 0
|
||||
|
||||
print('Load succeeded')
|
||||
|
||||
def super_sampling(self, step=1):
|
||||
with torch.no_grad():
|
||||
self.forward(step)
|
||||
self.fps = self.args.data.fps / step
|
||||
self.compute_test_result()
|
||||
|
||||
def compute_test_result(self):
|
||||
os.makedirs(self.viz_dir, exist_ok=True)
|
||||
|
||||
b_size = self.input_data['pos'].shape[0]
|
||||
n_joints = self.input_data['pos'].shape[2]
|
||||
|
||||
rotmat = self.recon_data['rotmat'] # (B, T, J, 3, 3)
|
||||
local_rotmat = self.fk.global_to_local(rotmat.view(-1, n_joints, 3, 3)) # (B x T, J, 3, 3)
|
||||
local_rotmat = local_rotmat.view(b_size, -1, n_joints, 3, 3) # (B, T, J, 3, 3)
|
||||
rotmat_gt = self.input_data['rotmat'] # (B, T, J, 3, 3)
|
||||
local_rotmat_gt = self.fk.global_to_local(rotmat_gt.view(-1, n_joints, 3, 3)) # (B x T, J, 3, 3)
|
||||
local_rotmat_gt = local_rotmat_gt.view(b_size, -1, n_joints, 3, 3) # (B, T, J, 3, 3)
|
||||
|
||||
if self.args.data.root_transform:
|
||||
root_orient = rotation_6d_to_matrix(self.recon_data['root_orient']) # (B, T, 3, 3)
|
||||
root_orient_gt = rotation_6d_to_matrix(self.input_data['root_orient']) # (B, T, 3, 3)
|
||||
local_rotmat[:, :, 0] = root_orient
|
||||
local_rotmat_gt[:, :, 0] = root_orient_gt
|
||||
rotations = matrix_to_quaternion(local_rotmat) # (B, T, J, 4)
|
||||
rotations_gt = matrix_to_quaternion(local_rotmat_gt) # (B, T, J, 4)
|
||||
|
||||
if 'trans' in self.recon_data.keys():
|
||||
origin = self.input_data['trans'][:, 0]
|
||||
positions = self.recon_data['trans']
|
||||
positions[..., self.v_axis] = positions[..., self.v_axis] + origin[..., self.v_axis].unsqueeze(1)
|
||||
positions_gt = self.input_data['trans'] # (B, T, 3)
|
||||
else:
|
||||
positions = torch.zeros_like(self.input_data['trans']) # (B, T, 3)
|
||||
positions_gt = torch.zeros_like(self.input_data['trans']) # (B, T, 3)
|
||||
|
||||
for i in range(b_size):
|
||||
if self.args.bvh_viz:
|
||||
rotation = align_joints(rotations[i])
|
||||
position = positions[i].unsqueeze(1) # (T, 1, 3)
|
||||
anim = Animation(Quaternions(c2c(rotation)), c2c(position), None, offsets=self.args.smpl.offsets[self.args.data.gender], parents=self.args.smpl.parents)
|
||||
BVH.save(os.path.join(self.viz_dir, f'test_{self.test_index + i:03d}_{int(self.fps)}fps.bvh'), anim, names=self.args.smpl.joint_names, frametime=1 / self.fps)
|
||||
|
||||
rotation_gt = align_joints(rotations_gt[i])
|
||||
position_gt = positions_gt[i].unsqueeze(1) # (T, 1, 3)
|
||||
anim.rotations = Quaternions(c2c(rotation_gt))
|
||||
anim.positions = c2c(position_gt)
|
||||
BVH.save(os.path.join(self.viz_dir, f'test_{self.test_index + i:03d}_gt.bvh'), anim, names=self.args.smpl.joint_names, frametime=1 / self.args.data.fps)
|
||||
else:
|
||||
poses = c2c(quaternion_to_axis_angle(rotations[i])) # (T, J, 3)
|
||||
poses_gt = c2c(quaternion_to_axis_angle(rotations_gt[i])) # (T, J, 3)
|
||||
|
||||
poses = poses.reshape((poses.shape[0], -1)) # (T, J x 3)
|
||||
poses = np.pad(poses, [(0, 0), (0, 93)], mode='constant')
|
||||
poses_gt = poses_gt.reshape((poses_gt.shape[0], -1)) # (T, J x 3)
|
||||
poses_gt = np.pad(poses_gt, [(0, 0), (0, 93)], mode='constant')
|
||||
|
||||
np.savez(os.path.join(self.viz_dir, f'test_{self.test_index + i:03d}_{int(self.fps)}fps.npz'),
|
||||
poses=poses, trans=c2c(positions[i]), betas=np.zeros(10), gender=self.args.data.gender, mocap_framerate=self.fps)
|
||||
np.savez(os.path.join(self.viz_dir, f'test_{self.test_index + i:03d}_gt.npz'),
|
||||
poses=poses_gt, trans=c2c(positions_gt[i]), betas=np.zeros(10), gender=self.args.data.gender, mocap_framerate=self.args.data.fps)
|
||||
|
||||
self.test_index += b_size
|
||||
@@ -0,0 +1,309 @@
|
||||
import os
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from human_body_prior.tools.omni_tools import copy2cpu as c2c
|
||||
from rotations import matrix_to_axis_angle, rotation_6d_to_matrix
|
||||
from utils import compute_trajectory, export_ply_trajectory, normalize
|
||||
|
||||
from .base_model import BaseModel
|
||||
from .fk import ForwardKinematicsLayer
|
||||
from .residual_blocks import ResidualBlock, SkeletonResidual, residual_ratio
|
||||
from .skeleton import SkeletonConv, SkeletonPool, build_edge_topology, find_neighbor
|
||||
|
||||
|
||||
class Predictor(nn.Module):
|
||||
def __init__(self, args, topology):
|
||||
super(Predictor, self).__init__()
|
||||
self.topologies = [topology]
|
||||
self.channel_base = [args.channel_base]
|
||||
self.channel_list = []
|
||||
self.edge_num = [len(topology)]
|
||||
self.pooling_list = []
|
||||
self.layers = nn.ModuleList()
|
||||
self.args = args
|
||||
# self.convs = []
|
||||
|
||||
kernel_size = args.kernel_size
|
||||
padding = (kernel_size - 1) // 2
|
||||
bias = True
|
||||
|
||||
for _ in range(args.num_layers):
|
||||
self.channel_base.append(self.channel_base[-1] * 2)
|
||||
|
||||
for i in range(args.num_layers):
|
||||
seq = []
|
||||
neighbour_list = find_neighbor(self.topologies[i], args.skeleton_dist)
|
||||
in_channels = self.channel_base[i] * self.edge_num[i]
|
||||
out_channels = self.channel_base[i + 1] * self.edge_num[i]
|
||||
if i == 0:
|
||||
self.channel_list.append(in_channels)
|
||||
self.channel_list.append(out_channels)
|
||||
last_pool = True if i == args.num_layers - 1 else False
|
||||
|
||||
# (T, J, D) => (T, J', D)
|
||||
pool = SkeletonPool(edges=self.topologies[i], pooling_mode=args.skeleton_pool,
|
||||
channels_per_edge=out_channels // len(neighbour_list), last_pool=last_pool)
|
||||
|
||||
if args.use_residual_blocks:
|
||||
# (T, J, D) => (T, J', 2D)
|
||||
seq.append(SkeletonResidual(self.topologies[i], neighbour_list, joint_num=self.edge_num[i], in_channels=in_channels, out_channels=out_channels,
|
||||
kernel_size=kernel_size, stride=1, padding=padding, padding_mode=args.padding_mode, bias=bias,
|
||||
extra_conv=args.extra_conv, pooling_mode=args.skeleton_pool, activation=args.activation, last_pool=last_pool))
|
||||
else:
|
||||
for _ in range(args.extra_conv):
|
||||
# (T, J, D) => (T, J, D)
|
||||
seq.append(SkeletonConv(neighbour_list, in_channels=in_channels, out_channels=in_channels,
|
||||
joint_num=self.edge_num[i], kernel_size=kernel_size, stride=1,
|
||||
padding=padding, padding_mode=args.padding_mode, bias=bias))
|
||||
seq.append(nn.PReLU())
|
||||
# (T, J, D) => (T, J, 2D)
|
||||
seq.append(SkeletonConv(neighbour_list, in_channels=in_channels, out_channels=out_channels,
|
||||
joint_num=self.edge_num[i], kernel_size=kernel_size, stride=1,
|
||||
padding=padding, padding_mode=args.padding_mode, bias=bias, add_offset=False,
|
||||
in_offset_channel=3 * self.channel_base[i] // self.channel_base[0]))
|
||||
# self.convs.append(seq[-1])
|
||||
|
||||
seq.append(pool)
|
||||
seq.append(nn.PReLU())
|
||||
self.layers.append(nn.Sequential(*seq))
|
||||
|
||||
self.topologies.append(pool.new_edges)
|
||||
self.pooling_list.append(pool.pooling_list)
|
||||
self.edge_num.append(len(self.topologies[-1]))
|
||||
|
||||
in_channels = self.channel_base[-1] * len(self.pooling_list[-1])
|
||||
out_channels = args.out_channels # root orient (6) + root vel (3) + root height (1) + contacts (8)
|
||||
|
||||
self.global_motion = nn.Sequential(
|
||||
ResidualBlock(in_channels=in_channels, out_channels=512, kernel_size=kernel_size, stride=1, padding=padding, residual_ratio=residual_ratio(1), activation=args.activation),
|
||||
ResidualBlock(in_channels=512, out_channels=256, kernel_size=kernel_size, stride=1, padding=padding, residual_ratio=residual_ratio(2), activation=args.activation),
|
||||
ResidualBlock(in_channels=256, out_channels=128, kernel_size=kernel_size, stride=1, padding=padding, residual_ratio=residual_ratio(3), activation=args.activation),
|
||||
ResidualBlock(in_channels=128, out_channels=out_channels, kernel_size=kernel_size, stride=1, padding=padding,
|
||||
residual_ratio=residual_ratio(4), activation=args.activation, last_layer=True)
|
||||
)
|
||||
|
||||
def forward(self, input):
|
||||
feature = input
|
||||
for layer in self.layers:
|
||||
feature = layer(feature)
|
||||
|
||||
return self.global_motion(feature)
|
||||
|
||||
|
||||
class GlobalMotionPredictor(BaseModel):
|
||||
def __init__(self, args, ngpu):
|
||||
super(GlobalMotionPredictor, self).__init__(args)
|
||||
|
||||
self.args = args
|
||||
# smpl_fname = os.path.join(args.smpl.smpl_body_model, args.data.gender, 'model.npz')
|
||||
# smpl_data = np.load(smpl_fname, encoding='latin1')
|
||||
smpl_fname = os.path.join(args.smpl.smpl_body_model, "MANO_RIGHT.pkl")
|
||||
smpl_data = np.load(smpl_fname, encoding='latin1', allow_pickle=True)
|
||||
parents = smpl_data['kintree_table'][0].astype(np.int32)
|
||||
edges = build_edge_topology(parents)
|
||||
|
||||
self.fk = ForwardKinematicsLayer(args)
|
||||
|
||||
self.predictor = Predictor(args.global_motion, edges).to(self.device)
|
||||
if args.multi_gpu is True:
|
||||
self.predictor = nn.DataParallel(self.predictor, device_ids=range(ngpu))
|
||||
self.models = [self.predictor]
|
||||
|
||||
ordermap = {
|
||||
'x': 0,
|
||||
'y': 1,
|
||||
'z': 2,
|
||||
}
|
||||
self.up_axis = ordermap[self.args.data.up]
|
||||
self.v_axis = [x for x in ordermap.values() if x != self.up_axis]
|
||||
|
||||
print(self.up_axis)
|
||||
print(self.v_axis)
|
||||
|
||||
self.input_data = {}
|
||||
self.pred_data = {}
|
||||
|
||||
self.mean = torch.load(os.path.join(args.dataset_dir, f'mean-{args.data.gender}-{args.data.clip_length}-{args.data.fps}fps.pt'), map_location=self.device)
|
||||
print(f'mean: {list(self.mean.keys())}')
|
||||
self.std = torch.load(os.path.join(args.dataset_dir, f'std-{args.data.gender}-{args.data.clip_length}-{args.data.fps}fps.pt'), map_location=self.device)
|
||||
print(f'std: {list(self.std.keys())}')
|
||||
|
||||
if self.is_train:
|
||||
self.optimizer = torch.optim.Adam(self.predictor.parameters(), args.learning_rate, weight_decay=args.weight_decay)
|
||||
self.optimizers = [self.optimizer]
|
||||
self.criterion_pred = nn.L1Loss() if args.l1_loss else nn.MSELoss()
|
||||
self.criterion_bce = nn.BCEWithLogitsLoss()
|
||||
else:
|
||||
self.test_index = 0
|
||||
self.smpl_dir = os.path.join(args.save_dir, 'results', 'smpl')
|
||||
|
||||
def set_input(self, input):
|
||||
self.input_data = {k: v.to(self.device) for k, v in input.items() if k in ['pos', 'velocity', 'rot6d', 'angular', 'global_xform',
|
||||
'root_vel', 'trans', 'contacts']}
|
||||
|
||||
def forward(self):
|
||||
self.pred_data.clear()
|
||||
self.input_data['origin'] = self.input_data['trans'][:, 0]
|
||||
self.pred_data = self.predict(data=self.input_data, dt=1.0 / self.args.data.fps)
|
||||
|
||||
def predict(self, data, dt):
|
||||
b_size, t_length = data['pos'].shape[:2]
|
||||
|
||||
if 'pos' in self.args.data.normalize:
|
||||
pos = normalize(data['pos'], mean=self.mean['pos'], std=self.std['pos'])
|
||||
else:
|
||||
pos = data['pos']
|
||||
if 'velocity' in self.args.data.normalize:
|
||||
velocity = normalize(data['velocity'], mean=self.mean['velocity'], std=self.std['velocity'])
|
||||
else:
|
||||
velocity = data['velocity']
|
||||
if 'rot6d' in self.args.data.normalize:
|
||||
rot6d = normalize(data['rot6d'], mean=self.mean['rot6d'], std=self.std['rot6d'])
|
||||
else:
|
||||
rot6d = data['rot6d']
|
||||
if 'angular' in self.args.data.normalize:
|
||||
angular = normalize(data['angular'], mean=self.mean['angular'], std=self.std['angular'])
|
||||
else:
|
||||
angular = data['angular']
|
||||
|
||||
x = torch.cat((pos, velocity, rot6d, angular), dim=-1)
|
||||
x = x.view(b_size, t_length, -1) # (B, T, J x D)
|
||||
x = x.permute(0, 2, 1) # (B, J x D, T)
|
||||
|
||||
global_motion = self.predictor(x)
|
||||
global_motion = global_motion.permute(0, 2, 1) # (B, T, D)
|
||||
|
||||
output = dict()
|
||||
output['root_vel'] = global_motion[:, :, :3]
|
||||
output['root_height'] = global_motion[:, :, 3]
|
||||
output['contacts'] = global_motion[:, :, 4:]
|
||||
if 'origin' in data.keys():
|
||||
origin = data['origin']
|
||||
else:
|
||||
origin = torch.zeros(b_size, 3).to(self.device)
|
||||
output['trans'] = compute_trajectory(output['root_vel'], output['root_height'], origin, dt, up_axis=self.args.data.up)
|
||||
|
||||
return output
|
||||
|
||||
def backward(self, validation=False):
|
||||
loss = 0
|
||||
|
||||
v_loss = self.criterion_pred(self.pred_data['root_vel'], self.input_data['root_vel'])
|
||||
up_loss = self.criterion_pred(self.pred_data['root_height'], self.input_data['trans'][..., self.up_axis])
|
||||
self.loss_recorder.add_scalar('v_loss', v_loss, validation=validation)
|
||||
self.loss_recorder.add_scalar('up_loss', up_loss, validation=validation)
|
||||
loss += self.args.lambda_v * v_loss + self.args.lambda_up * up_loss
|
||||
|
||||
trans_loss = self.criterion_pred(self.pred_data['trans'], self.input_data['trans'])
|
||||
self.loss_recorder.add_scalar('trans_loss', trans_loss, validation=validation)
|
||||
loss += self.args.lambda_trans * trans_loss
|
||||
|
||||
contacts_loss = self.criterion_bce(self.pred_data['contacts'], self.input_data['contacts'])
|
||||
self.loss_recorder.add_scalar('contacts_loss', contacts_loss, validation=validation)
|
||||
loss += self.args.lambda_contacts * contacts_loss
|
||||
|
||||
self.loss_recorder.add_scalar('total_loss', loss, validation=validation)
|
||||
|
||||
if not validation:
|
||||
loss.backward()
|
||||
|
||||
def optimize_parameters(self):
|
||||
self.optimizer.zero_grad()
|
||||
self.forward()
|
||||
self.backward(validation=False)
|
||||
self.optimizer.step()
|
||||
|
||||
def report_errors(self):
|
||||
trans = self.pred_data['trans'] # (B, T, 3)
|
||||
trans_gt = self.input_data['trans'] # (B, T, 3)
|
||||
translation_error = torch.linalg.norm((trans - trans_gt), dim=-1).mean()
|
||||
|
||||
return {
|
||||
'translation': c2c(translation_error)
|
||||
}
|
||||
|
||||
def verbose(self):
|
||||
res = {}
|
||||
for loss in self.loss_recorder.losses.values():
|
||||
res[loss.name] = {'train': loss.current()[0], 'val': loss.current()[1]}
|
||||
|
||||
return res
|
||||
|
||||
def validate(self):
|
||||
with torch.no_grad():
|
||||
self.forward()
|
||||
self.backward(validation=True)
|
||||
|
||||
def save(self, optimal=False):
|
||||
if optimal:
|
||||
path = os.path.join(self.args.save_dir, 'results', 'model')
|
||||
else:
|
||||
path = os.path.join(self.model_save_dir, f'{self.epoch_cnt:04d}')
|
||||
|
||||
os.makedirs(path, exist_ok=True)
|
||||
predictor = self.predictor.module if isinstance(self.predictor, nn.DataParallel) else self.predictor
|
||||
torch.save(predictor.state_dict(), os.path.join(path, 'predictor.pth'))
|
||||
torch.save(self.optimizer.state_dict(), os.path.join(path, 'optimizer.pth'))
|
||||
if self.args.scheduler.name:
|
||||
torch.save(self.schedulers[0].state_dict(), os.path.join(path, 'scheduler.pth'))
|
||||
self.loss_recorder.save(path)
|
||||
|
||||
print(f'Save at {path} succeeded')
|
||||
|
||||
def load(self, epoch=None, optimal=False):
|
||||
if optimal:
|
||||
path = os.path.join(self.args.save_dir, 'results', 'model')
|
||||
else:
|
||||
if epoch is None:
|
||||
all = [int(q) for q in os.listdir(self.model_save_dir)]
|
||||
if len(all) == 0:
|
||||
raise RuntimeError(f'Empty loading path {self.model_save_dir}')
|
||||
epoch = sorted(all)[-1]
|
||||
path = os.path.join(self.model_save_dir, f'{epoch:04d}')
|
||||
|
||||
print(f'Loading from {path}')
|
||||
predictor = self.predictor.module if isinstance(self.predictor, nn.DataParallel) else self.predictor
|
||||
predictor.load_state_dict(torch.load(os.path.join(path, 'predictor.pth'), map_location=self.device))
|
||||
if self.is_train:
|
||||
self.optimizer.load_state_dict(torch.load(os.path.join(path, 'optimizer.pth')))
|
||||
self.loss_recorder.load(path)
|
||||
if self.args.scheduler.name:
|
||||
self.schedulers[0].load_state_dict(torch.load(os.path.join(path, 'scheduler.pth')))
|
||||
self.epoch_cnt = epoch if not optimal else 0
|
||||
|
||||
print('Load succeeded')
|
||||
|
||||
def compute_test_result(self):
|
||||
os.makedirs(self.smpl_dir, exist_ok=True)
|
||||
|
||||
b_size = self.input_data['trans'].shape[0]
|
||||
n_joints = self.input_data['pos'].shape[2]
|
||||
|
||||
trans = self.pred_data['trans']
|
||||
trans_gt = self.input_data['trans']
|
||||
|
||||
rotmat_gt = rotation_6d_to_matrix(self.input_data['global_xform']) # (B, T, J, 3, 3)
|
||||
local_rotmat_gt = self.fk.global_to_local(rotmat_gt.view(-1, n_joints, 3, 3)) # (B x T, J, 3, 3)
|
||||
local_rotmat_gt = local_rotmat_gt.view(b_size, -1, n_joints, 3, 3) # (B, T, J, 3, 3)
|
||||
local_rotmat = local_rotmat_gt.clone()
|
||||
|
||||
for i in range(b_size):
|
||||
export_ply_trajectory(points=trans[i], color=(255, 0, 0), ply_fname=os.path.join(self.smpl_dir, f'test_{self.test_index + i:03d}.ply'))
|
||||
export_ply_trajectory(points=trans_gt[i], color=(0, 255, 0), ply_fname=os.path.join(self.smpl_dir, f'test_{self.test_index + i:03d}_gt.ply'))
|
||||
|
||||
poses = c2c(matrix_to_axis_angle(local_rotmat[i])) # (T, J, 3)
|
||||
poses_gt = c2c(matrix_to_axis_angle(local_rotmat_gt[i])) # (T, J, 3)
|
||||
|
||||
poses = poses.reshape((poses.shape[0], -1)) # (T, J x 3)
|
||||
poses = np.pad(poses, [(0, 0), (0, 93)], mode='constant')
|
||||
poses_gt = poses_gt.reshape((poses_gt.shape[0], -1)) # (T, J x 3)
|
||||
poses_gt = np.pad(poses_gt, [(0, 0), (0, 93)], mode='constant')
|
||||
|
||||
np.savez(os.path.join(self.smpl_dir, f'test_{self.test_index + i:03d}.npz'),
|
||||
poses=poses, trans=c2c(trans[i]), betas=np.zeros(10), gender=self.args.data.gender, mocap_framerate=self.args.data.fps)
|
||||
np.savez(os.path.join(self.smpl_dir, f'test_{self.test_index + i:03d}_gt.npz'),
|
||||
poses=poses_gt, trans=c2c(trans_gt[i]), betas=np.zeros(10), gender=self.args.data.gender, mocap_framerate=self.args.data.fps)
|
||||
|
||||
self.test_index += b_size
|
||||
@@ -0,0 +1,82 @@
|
||||
import os
|
||||
import pickle
|
||||
|
||||
import torch
|
||||
from tensorboardX import SummaryWriter
|
||||
|
||||
|
||||
class SingleLoss:
|
||||
def __init__(self, name: str, writer: SummaryWriter):
|
||||
self.name = name
|
||||
self.loss_step = []
|
||||
self.loss_epoch = []
|
||||
self.loss_epoch_tmp = []
|
||||
self.loss_epoch_val = []
|
||||
self.writer = writer
|
||||
|
||||
def add_scalar(self, val, step=None, validation=False):
|
||||
if validation:
|
||||
self.loss_epoch_val.append(val)
|
||||
else:
|
||||
if step is None:
|
||||
step = len(self.loss_step)
|
||||
self.loss_step.append(val)
|
||||
self.loss_epoch_tmp.append(val)
|
||||
self.writer.add_scalar('Loss/step_' + self.name, val, step)
|
||||
|
||||
def current(self):
|
||||
return self.loss_epoch[-1]
|
||||
|
||||
def epoch(self, step=None):
|
||||
if step is None:
|
||||
step = len(self.loss_epoch)
|
||||
if self.loss_epoch_tmp:
|
||||
loss_avg = sum(self.loss_epoch_tmp) / len(self.loss_epoch_tmp)
|
||||
else:
|
||||
loss_avg = 0
|
||||
self.loss_epoch_tmp = []
|
||||
if self.loss_epoch_val:
|
||||
loss_avg_val = sum(self.loss_epoch_val) / len(self.loss_epoch_val)
|
||||
else:
|
||||
loss_avg_val = 0
|
||||
self.loss_epoch_val = []
|
||||
self.loss_epoch.append([loss_avg, loss_avg_val])
|
||||
self.writer.add_scalars('Loss/epoch_' + self.name, {'train': loss_avg, 'val': loss_avg_val}, step)
|
||||
|
||||
def save(self):
|
||||
return {
|
||||
'loss_step': self.loss_step,
|
||||
'loss_epoch': self.loss_epoch
|
||||
}
|
||||
|
||||
|
||||
class LossRecorder:
|
||||
def __init__(self, writer: SummaryWriter):
|
||||
self.losses = {}
|
||||
self.writer = writer
|
||||
|
||||
def add_scalar(self, name, val, step=None, validation=False):
|
||||
if isinstance(val, torch.Tensor):
|
||||
val = val.item()
|
||||
if name not in self.losses:
|
||||
self.losses[name] = SingleLoss(name, self.writer)
|
||||
self.losses[name].add_scalar(val, step, validation)
|
||||
|
||||
def epoch(self, step=None):
|
||||
for loss in self.losses.values():
|
||||
loss.epoch(step)
|
||||
|
||||
def save(self, path):
|
||||
data = {}
|
||||
for key, value in self.losses.items():
|
||||
data[key] = value.save()
|
||||
with open(os.path.join(path, 'loss.pkl'), 'wb') as f:
|
||||
pickle.dump(data, f)
|
||||
|
||||
def load(self, path):
|
||||
with open(os.path.join(path, 'loss.pkl'), 'rb') as f:
|
||||
data = pickle.load(f)
|
||||
for key, value in data.items():
|
||||
self.losses[key] = SingleLoss(key, self.writer)
|
||||
self.losses[key].loss_step = value['loss_step']
|
||||
self.losses[key].loss_epoch = value['loss_epoch']
|
||||
@@ -0,0 +1,88 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class GeodesicLoss(nn.Module):
|
||||
def __init__(self):
|
||||
super(GeodesicLoss, self).__init__()
|
||||
|
||||
def compute_geodesic_distance(self, m1, m2):
|
||||
""" Compute the geodesic distance between two rotation matrices.
|
||||
|
||||
Args:
|
||||
m1, m2: Two rotation matrices with the shape (batch x 3 x 3).
|
||||
|
||||
Returns:
|
||||
The minimal angular difference between two rotation matrices in radian form [0, pi].
|
||||
"""
|
||||
batch = m1.shape[0]
|
||||
m = torch.bmm(m1, m2.transpose(1, 2)) # batch*3*3
|
||||
|
||||
cos = (m[:, 0, 0] + m[:, 1, 1] + m[:, 2, 2] - 1) / 2
|
||||
# cos = torch.min(cos, torch.autograd.Variable(torch.ones(batch).cuda()))
|
||||
# cos = torch.max(cos, torch.autograd.Variable(torch.ones(batch).cuda()) * -1)
|
||||
eps = 1e-7
|
||||
cos = torch.clamp(cos, -1 + eps, 1 - eps)
|
||||
|
||||
theta = torch.acos(cos)
|
||||
|
||||
return theta
|
||||
|
||||
def rot_smooth_loss(self, rotmat):
|
||||
|
||||
diff = self.compute_geodesic_distance(rotmat[:, 1:], rotmat[:, :-1])
|
||||
vel = diff * 30
|
||||
loss = vel ** 2
|
||||
loss = 0.5 * torch.mean(loss)
|
||||
return loss
|
||||
|
||||
def __call__(self, m1, m2, reduction='mean'):
|
||||
loss = self.compute_geodesic_distance(m1, m2)
|
||||
|
||||
if reduction == 'mean':
|
||||
return loss.mean()
|
||||
elif reduction == 'none':
|
||||
return loss
|
||||
else:
|
||||
raise RuntimeError(f'unsupported reduction: {reduction}')
|
||||
|
||||
|
||||
def pos_smooth_loss(pos):
|
||||
# minimize delta steps
|
||||
diff = (pos[:, 1:] - pos[:, :-1])
|
||||
vel = diff * 30
|
||||
loss = vel ** 2
|
||||
loss = 0.5 * torch.mean(loss)
|
||||
return loss
|
||||
|
||||
def compute_geodesic_distance(m1, m2):
|
||||
""" Compute the geodesic distance between two rotation matrices.
|
||||
|
||||
Args:
|
||||
m1, m2: Two rotation matrices with the shape (batch x 3 x 3).
|
||||
|
||||
Returns:
|
||||
The minimal angular difference between two rotation matrices in radian form [0, pi].
|
||||
"""
|
||||
batch = m1.shape[:-2]
|
||||
m = torch.matmul(m1, m2.transpose(-2, -1)) # batch*3*3
|
||||
|
||||
cos = (m[..., 0, 0] + m[..., 1, 1] + m[..., 2, 2] - 1) / 2
|
||||
cos = torch.min(cos, torch.ones(*batch).cuda())
|
||||
cos = torch.max(cos, torch.ones(*batch).cuda() * -1)
|
||||
|
||||
eps = 1e-7
|
||||
cos = torch.clamp(cos, -1 + eps, 1 - eps)
|
||||
|
||||
theta = torch.acos(cos)
|
||||
return theta
|
||||
|
||||
|
||||
def rot_smooth_loss(rotmat):
|
||||
|
||||
diff = compute_geodesic_distance(rotmat[:, 1:], rotmat[:, :-1])
|
||||
vel = diff * 30
|
||||
loss = vel ** 2
|
||||
loss = 0.5 * torch.mean(loss)
|
||||
return loss
|
||||
|
||||
@@ -0,0 +1,161 @@
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
|
||||
class PositionalEncoding(nn.Module):
|
||||
def __init__(self, n_functions):
|
||||
super(PositionalEncoding, self).__init__()
|
||||
|
||||
self.register_buffer('frequencies', 2.0 ** torch.arange(n_functions))
|
||||
|
||||
def forward(self, x):
|
||||
"""
|
||||
Args:
|
||||
x: tensor of shape [..., dim]
|
||||
|
||||
Returns:
|
||||
embedding: a temporal embedding of `x` of shape [..., n_functions * dim * 2]
|
||||
"""
|
||||
freq = (x[..., None] * self.frequencies).view(*x.shape[:-1], -1)
|
||||
|
||||
embedding = torch.zeros(*freq.shape[:-1], freq.shape[-1] * 2).to(device)
|
||||
embedding[..., 0::2] = freq.sin()
|
||||
embedding[..., 1::2] = freq.cos()
|
||||
|
||||
return embedding
|
||||
|
||||
|
||||
class Sine(nn.Module):
|
||||
def __init__(self, w0=1.):
|
||||
super(Sine, self).__init__()
|
||||
|
||||
self.w0 = w0
|
||||
|
||||
def forward(self, x):
|
||||
return torch.sin(self.w0 * x)
|
||||
|
||||
|
||||
class SirenBlock(nn.Module):
|
||||
def __init__(self, in_features, out_features, w0=30, c=6, is_first=False, use_bias=True, activation=None):
|
||||
super(SirenBlock, self).__init__()
|
||||
|
||||
self.in_features = in_features
|
||||
self.is_first = is_first
|
||||
|
||||
weight = torch.zeros(out_features, in_features)
|
||||
bias = torch.zeros(out_features) if use_bias else None
|
||||
self.init(weight, bias, c=c, w0=w0)
|
||||
|
||||
self.weight = nn.Parameter(weight)
|
||||
self.bias = nn.Parameter(bias) if use_bias else None
|
||||
self.activation = Sine(w0) if activation is None else activation
|
||||
|
||||
def init(self, weight, bias, c, w0):
|
||||
n = self.in_features
|
||||
|
||||
w_std = (1 / n) if self.is_first else (np.sqrt(c / n) / w0)
|
||||
weight.uniform_(-w_std, w_std)
|
||||
|
||||
if bias is not None:
|
||||
bias.uniform_(-w_std, w_std)
|
||||
|
||||
def forward(self, x):
|
||||
out = F.linear(x, self.weight, self.bias)
|
||||
|
||||
return self.activation(out)
|
||||
|
||||
|
||||
class FCBlock(nn.Module):
|
||||
def __init__(self, in_features, out_features, norm_layer=False, activation=None):
|
||||
super(FCBlock, self).__init__()
|
||||
|
||||
self.fc = nn.Linear(in_features, out_features)
|
||||
self.residual = (in_features == out_features) # when the input and output have the same dimensions, build a residual block
|
||||
self.norm_layer = nn.LayerNorm(out_features) if norm_layer else None
|
||||
self.activation = nn.ReLU(inplace=True) if activation is None else activation
|
||||
|
||||
def forward(self, x):
|
||||
"""
|
||||
Args:
|
||||
x: (B, T, D), features are in the last dimension.
|
||||
"""
|
||||
out = self.fc(x)
|
||||
|
||||
if self.norm_layer is not None:
|
||||
out = self.norm_layer(out)
|
||||
|
||||
if self.residual:
|
||||
return self.activation(out) + x
|
||||
|
||||
return self.activation(out)
|
||||
|
||||
|
||||
class NeuralMotionField(nn.Module):
|
||||
def __init__(self, args):
|
||||
super(NeuralMotionField, self).__init__()
|
||||
|
||||
self.args = args
|
||||
|
||||
if args.bandwidth != 0:
|
||||
self.positional_encoding = PositionalEncoding(args.bandwidth)
|
||||
embedding_dim = args.bandwidth * 2
|
||||
else:
|
||||
embedding_dim = 1
|
||||
|
||||
hidden_neuron = args.hidden_neuron
|
||||
in_features = 1 + args.local_z if args.siren else embedding_dim + args.local_z
|
||||
local_in_features = hidden_neuron + in_features if args.skip_connection else hidden_neuron
|
||||
global_in_features = local_in_features + args.global_z if args.skip_connection else hidden_neuron
|
||||
|
||||
layers = [
|
||||
SirenBlock(in_features, hidden_neuron, is_first=True) if args.siren else FCBlock(in_features, hidden_neuron, norm_layer=args.norm_layer),
|
||||
SirenBlock(local_in_features, hidden_neuron) if args.siren else FCBlock(local_in_features, hidden_neuron, norm_layer=args.norm_layer),
|
||||
SirenBlock(local_in_features, hidden_neuron) if args.siren else FCBlock(local_in_features, hidden_neuron, norm_layer=args.norm_layer),
|
||||
SirenBlock(local_in_features, hidden_neuron) if args.siren else FCBlock(local_in_features, hidden_neuron, norm_layer=args.norm_layer),
|
||||
SirenBlock(local_in_features, hidden_neuron) if args.siren else FCBlock(local_in_features, hidden_neuron, norm_layer=args.norm_layer),
|
||||
SirenBlock(local_in_features, hidden_neuron) if args.siren else FCBlock(local_in_features, hidden_neuron, norm_layer=args.norm_layer),
|
||||
SirenBlock(local_in_features, hidden_neuron) if args.siren else FCBlock(local_in_features, hidden_neuron, norm_layer=args.norm_layer),
|
||||
SirenBlock(local_in_features, hidden_neuron) if args.siren else FCBlock(local_in_features, hidden_neuron, norm_layer=args.norm_layer),
|
||||
SirenBlock(global_in_features, hidden_neuron) if args.siren else FCBlock(global_in_features, hidden_neuron, norm_layer=args.norm_layer),
|
||||
SirenBlock(global_in_features, hidden_neuron) if args.siren else FCBlock(global_in_features, hidden_neuron, norm_layer=args.norm_layer),
|
||||
SirenBlock(global_in_features, hidden_neuron) if args.siren else FCBlock(global_in_features, hidden_neuron, norm_layer=args.norm_layer)
|
||||
]
|
||||
self.mlp = nn.ModuleList(layers)
|
||||
self.skip_layers = [] if not args.skip_connection else list(range(1, len(self.mlp)))
|
||||
self.local_layers = list(range(8))
|
||||
self.local_linear = nn.Sequential(nn.Linear(hidden_neuron, args.local_output))
|
||||
self.global_layers = list(range(8, len(self.mlp)))
|
||||
self.global_linear = nn.Sequential(nn.Linear(hidden_neuron, args.global_output))
|
||||
|
||||
def forward(self, t, z_l=None, z_g=None):
|
||||
if self.args.bandwidth != 0 and self.args.siren is False:
|
||||
t = self.positional_encoding(t)
|
||||
|
||||
if z_l is not None:
|
||||
z_l = z_l.unsqueeze(1).expand(-1, t.shape[1], -1)
|
||||
x = torch.cat([t, z_l], dim=-1)
|
||||
else:
|
||||
x = t
|
||||
|
||||
skip_in = x
|
||||
for i in self.local_layers:
|
||||
if i in self.skip_layers:
|
||||
x = torch.cat((x, skip_in), dim=-1)
|
||||
x = self.mlp[i](x)
|
||||
local_output = self.local_linear(x)
|
||||
|
||||
if z_g is not None:
|
||||
z_g = z_g.unsqueeze(1).expand(-1, t.shape[1], -1)
|
||||
for i in self.global_layers:
|
||||
if i in self.skip_layers:
|
||||
x = torch.cat((x, skip_in, z_g), dim=-1)
|
||||
x = self.mlp[i](x)
|
||||
global_output = self.global_linear(x)
|
||||
else:
|
||||
global_output = None
|
||||
|
||||
return local_output, global_output
|
||||
+109
@@ -0,0 +1,109 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from .residual_blocks import ResidualBlock, SkeletonResidual, residual_ratio
|
||||
from .skeleton import SkeletonConv, SkeletonPool, find_neighbor
|
||||
|
||||
|
||||
class LocalEncoder(nn.Module):
|
||||
def __init__(self, args, topology):
|
||||
super(LocalEncoder, self).__init__()
|
||||
self.topologies = [topology]
|
||||
self.channel_base = [args.channel_base]
|
||||
|
||||
self.channel_list = []
|
||||
self.edge_num = [len(topology)]
|
||||
self.pooling_list = []
|
||||
self.layers = nn.ModuleList()
|
||||
self.args = args
|
||||
# self.convs = []
|
||||
|
||||
kernel_size = args.kernel_size
|
||||
kernel_even = False if kernel_size % 2 else True
|
||||
padding = (kernel_size - 1) // 2
|
||||
bias = True
|
||||
|
||||
for _ in range(args.num_layers):
|
||||
self.channel_base.append(self.channel_base[-1] * 2)
|
||||
|
||||
for i in range(args.num_layers):
|
||||
seq = []
|
||||
neighbour_list = find_neighbor(self.topologies[i], args.skeleton_dist)
|
||||
in_channels = self.channel_base[i] * self.edge_num[i]
|
||||
out_channels = self.channel_base[i + 1] * self.edge_num[i]
|
||||
if i == 0:
|
||||
self.channel_list.append(in_channels)
|
||||
self.channel_list.append(out_channels)
|
||||
last_pool = True if i == args.num_layers - 1 else False
|
||||
|
||||
# (T, J, D) => (T, J', D)
|
||||
pool = SkeletonPool(edges=self.topologies[i], pooling_mode=args.skeleton_pool,
|
||||
channels_per_edge=out_channels // len(neighbour_list), last_pool=last_pool)
|
||||
|
||||
if args.use_residual_blocks:
|
||||
# (T, J, D) => (T/2, J', 2D)
|
||||
seq.append(SkeletonResidual(self.topologies[i], neighbour_list, joint_num=self.edge_num[i], in_channels=in_channels, out_channels=out_channels,
|
||||
kernel_size=kernel_size, stride=2, padding=padding, padding_mode=args.padding_mode, bias=bias,
|
||||
extra_conv=args.extra_conv, pooling_mode=args.skeleton_pool, activation=args.activation, last_pool=last_pool))
|
||||
else:
|
||||
for _ in range(args.extra_conv):
|
||||
# (T, J, D) => (T, J, D)
|
||||
seq.append(SkeletonConv(neighbour_list, in_channels=in_channels, out_channels=in_channels,
|
||||
joint_num=self.edge_num[i], kernel_size=kernel_size - 1 if kernel_even else kernel_size,
|
||||
stride=1,
|
||||
padding=padding, padding_mode=args.padding_mode, bias=bias))
|
||||
seq.append(nn.PReLU() if args.activation == 'relu' else nn.Tanh())
|
||||
# (T, J, D) => (T/2, J, 2D)
|
||||
seq.append(SkeletonConv(neighbour_list, in_channels=in_channels, out_channels=out_channels,
|
||||
joint_num=self.edge_num[i], kernel_size=kernel_size, stride=2,
|
||||
padding=padding, padding_mode=args.padding_mode, bias=bias, add_offset=False,
|
||||
in_offset_channel=3 * self.channel_base[i] // self.channel_base[0]))
|
||||
# self.convs.append(seq[-1])
|
||||
|
||||
seq.append(pool)
|
||||
seq.append(nn.PReLU() if args.activation == 'relu' else nn.Tanh())
|
||||
self.layers.append(nn.Sequential(*seq))
|
||||
|
||||
self.topologies.append(pool.new_edges)
|
||||
self.pooling_list.append(pool.pooling_list)
|
||||
self.edge_num.append(len(self.topologies[-1]))
|
||||
|
||||
in_features = self.channel_base[-1] * len(self.pooling_list[-1])
|
||||
in_features *= args.temporal_scale
|
||||
self.mu = nn.Linear(in_features, args.z_dim)
|
||||
self.logvar = nn.Linear(in_features, args.z_dim)
|
||||
|
||||
def forward(self, input):
|
||||
output = input
|
||||
for layer in self.layers:
|
||||
output = layer(output)
|
||||
output = output.view(output.shape[0], -1)
|
||||
|
||||
return self.mu(output), self.logvar(output)
|
||||
|
||||
|
||||
class GlobalEncoder(nn.Module):
|
||||
def __init__(self, args):
|
||||
super(GlobalEncoder, self).__init__()
|
||||
|
||||
in_channels = args.in_channels # root orientation + root velocity
|
||||
out_channels = 512
|
||||
kernel_size = args.kernel_size
|
||||
padding = (kernel_size - 1) // 2
|
||||
|
||||
self.encoder = nn.Sequential(
|
||||
ResidualBlock(in_channels=in_channels, out_channels=128, kernel_size=kernel_size, stride=2, padding=padding, residual_ratio=residual_ratio(1), activation=args.activation),
|
||||
ResidualBlock(in_channels=128, out_channels=256, kernel_size=kernel_size, stride=2, padding=padding, residual_ratio=residual_ratio(2), activation=args.activation),
|
||||
ResidualBlock(in_channels=256, out_channels=512, kernel_size=kernel_size, stride=2, padding=padding, residual_ratio=residual_ratio(3), activation=args.activation),
|
||||
ResidualBlock(in_channels=512, out_channels=out_channels, kernel_size=kernel_size, stride=2, padding=padding, residual_ratio=residual_ratio(4), activation=args.activation)
|
||||
)
|
||||
|
||||
in_features = out_channels * args.temporal_scale
|
||||
self.mu = nn.Linear(in_features, args.z_dim)
|
||||
self.logvar = nn.Linear(in_features, args.z_dim)
|
||||
|
||||
def forward(self, input):
|
||||
feature = self.encoder(input)
|
||||
feature = feature.view(feature.shape[0], -1)
|
||||
|
||||
return self.mu(feature), self.logvar(feature)
|
||||
@@ -0,0 +1,197 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from .skeleton import SkeletonConv, SkeletonPool, SkeletonUnpool
|
||||
|
||||
|
||||
def residual_ratio(k):
|
||||
return 1 / (k + 1)
|
||||
|
||||
|
||||
class Affine(nn.Module):
|
||||
def __init__(self, num_parameters, scale=True, bias=True, scale_init=1.0):
|
||||
super(Affine, self).__init__()
|
||||
if scale:
|
||||
self.scale = nn.Parameter(torch.ones(num_parameters) * scale_init)
|
||||
else:
|
||||
self.register_parameter('scale', None)
|
||||
|
||||
if bias:
|
||||
self.bias = nn.Parameter(torch.zeros(num_parameters))
|
||||
else:
|
||||
self.register_parameter('bias', None)
|
||||
|
||||
def forward(self, input):
|
||||
output = input
|
||||
if self.scale is not None:
|
||||
scale = self.scale.unsqueeze(0)
|
||||
while scale.dim() < input.dim():
|
||||
scale = scale.unsqueeze(2)
|
||||
output = output.mul(scale)
|
||||
|
||||
if self.bias is not None:
|
||||
bias = self.bias.unsqueeze(0)
|
||||
while bias.dim() < input.dim():
|
||||
bias = bias.unsqueeze(2)
|
||||
output += bias
|
||||
|
||||
return output
|
||||
|
||||
|
||||
class BatchStatistics(nn.Module):
|
||||
def __init__(self, affine=-1):
|
||||
super(BatchStatistics, self).__init__()
|
||||
self.affine = nn.Sequential() if affine == -1 else Affine(affine)
|
||||
self.loss = 0
|
||||
|
||||
def clear_loss(self):
|
||||
self.loss = 0
|
||||
|
||||
def compute_loss(self, input):
|
||||
input_flat = input.view(input.size(1), input.numel() // input.size(1))
|
||||
mu = input_flat.mean(1)
|
||||
logvar = (input_flat.pow(2).mean(1) - mu.pow(2)).sqrt().log()
|
||||
|
||||
self.loss = mu.pow(2).mean() + logvar.pow(2).mean()
|
||||
|
||||
def forward(self, input):
|
||||
self.compute_loss(input)
|
||||
return self.affine(input)
|
||||
|
||||
|
||||
class ResidualBlock(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, kernel_size, stride, padding, residual_ratio, activation, batch_statistics=False, last_layer=False):
|
||||
super(ResidualBlock, self).__init__()
|
||||
|
||||
self.residual_ratio = residual_ratio
|
||||
self.shortcut_ratio = 1 - residual_ratio
|
||||
|
||||
residual = []
|
||||
residual.append(nn.Conv1d(in_channels, out_channels, kernel_size, stride, padding))
|
||||
if batch_statistics:
|
||||
residual.append(BatchStatistics(out_channels))
|
||||
if not last_layer:
|
||||
residual.append(nn.PReLU() if activation == 'relu' else nn.Tanh())
|
||||
self.residual = nn.Sequential(*residual)
|
||||
|
||||
self.shortcut = nn.Sequential(
|
||||
nn.AvgPool1d(kernel_size=2) if stride == 2 else nn.Sequential(),
|
||||
nn.Conv1d(in_channels, out_channels, kernel_size=1, stride=1, padding=0),
|
||||
BatchStatistics(out_channels) if (in_channels != out_channels and batch_statistics is True) else nn.Sequential()
|
||||
)
|
||||
|
||||
def forward(self, input):
|
||||
return self.residual(input).mul(self.residual_ratio) + self.shortcut(input).mul(self.shortcut_ratio)
|
||||
|
||||
|
||||
class ResidualBlockTranspose(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, kernel_size, stride, padding, residual_ratio, activation):
|
||||
super(ResidualBlockTranspose, self).__init__()
|
||||
|
||||
self.residual_ratio = residual_ratio
|
||||
self.shortcut_ratio = 1 - residual_ratio
|
||||
|
||||
self.residual = nn.Sequential(
|
||||
nn.ConvTranspose1d(in_channels, out_channels, kernel_size, stride, padding),
|
||||
nn.PReLU() if activation == 'relu' else nn.Tanh()
|
||||
)
|
||||
|
||||
self.shortcut = nn.Sequential(
|
||||
nn.Upsample(scale_factor=2, mode='linear', align_corners=False) if stride == 2 else nn.Sequential(),
|
||||
nn.Conv1d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
|
||||
)
|
||||
|
||||
def forward(self, input):
|
||||
return self.residual(input).mul(self.residual_ratio) + self.shortcut(input).mul(self.shortcut_ratio)
|
||||
|
||||
|
||||
class SkeletonResidual(nn.Module):
|
||||
def __init__(self, topology, neighbour_list, joint_num, in_channels, out_channels, kernel_size, stride, padding, padding_mode, bias, extra_conv, pooling_mode, activation, last_pool):
|
||||
super(SkeletonResidual, self).__init__()
|
||||
|
||||
kernel_even = False if kernel_size % 2 else True
|
||||
|
||||
seq = []
|
||||
for _ in range(extra_conv):
|
||||
# (T, J, D) => (T, J, D)
|
||||
seq.append(SkeletonConv(neighbour_list, in_channels=in_channels, out_channels=in_channels,
|
||||
joint_num=joint_num, kernel_size=kernel_size - 1 if kernel_even else kernel_size,
|
||||
stride=1,
|
||||
padding=padding, padding_mode=padding_mode, bias=bias))
|
||||
seq.append(nn.PReLU() if activation == 'relu' else nn.Tanh())
|
||||
# (T, J, D) => (T/2, J, 2D)
|
||||
seq.append(SkeletonConv(neighbour_list, in_channels=in_channels, out_channels=out_channels,
|
||||
joint_num=joint_num, kernel_size=kernel_size, stride=stride,
|
||||
padding=padding, padding_mode=padding_mode, bias=bias, add_offset=False))
|
||||
seq.append(nn.GroupNorm(4, out_channels)) # FIXME: REMEMBER TO CHANGE BACK !!!
|
||||
self.residual = nn.Sequential(*seq)
|
||||
|
||||
# (T, J, D) => (T/2, J, 2D)
|
||||
self.shortcut = SkeletonConv(neighbour_list, in_channels=in_channels, out_channels=out_channels,
|
||||
joint_num=joint_num, kernel_size=1, stride=stride, padding=0,
|
||||
bias=True, add_offset=False)
|
||||
|
||||
seq = []
|
||||
# (T/2, J, 2D) => (T/2, J', 2D)
|
||||
pool = SkeletonPool(edges=topology, pooling_mode=pooling_mode,
|
||||
channels_per_edge=out_channels // len(neighbour_list), last_pool=last_pool)
|
||||
if len(pool.pooling_list) != pool.edge_num:
|
||||
seq.append(pool)
|
||||
seq.append(nn.PReLU() if activation == 'relu' else nn.Tanh())
|
||||
self.common = nn.Sequential(*seq)
|
||||
|
||||
def forward(self, input):
|
||||
output = self.residual(input) + self.shortcut(input)
|
||||
|
||||
return self.common(output)
|
||||
|
||||
|
||||
class SkeletonResidualTranspose(nn.Module):
|
||||
def __init__(self, neighbour_list, joint_num, in_channels, out_channels, kernel_size, padding, padding_mode, bias, extra_conv, pooling_list, upsampling, activation, last_layer):
|
||||
super(SkeletonResidualTranspose, self).__init__()
|
||||
|
||||
kernel_even = False if kernel_size % 2 else True
|
||||
|
||||
seq = []
|
||||
# (T, J, D) => (2T, J, D)
|
||||
if upsampling is not None:
|
||||
seq.append(nn.Upsample(scale_factor=2, mode=upsampling, align_corners=False))
|
||||
# (2T, J, D) => (2T, J', D)
|
||||
unpool = SkeletonUnpool(pooling_list, in_channels // len(neighbour_list))
|
||||
if unpool.input_edge_num != unpool.output_edge_num:
|
||||
seq.append(unpool)
|
||||
self.common = nn.Sequential(*seq)
|
||||
|
||||
seq = []
|
||||
for _ in range(extra_conv):
|
||||
# (2T, J', D) => (2T, J', D)
|
||||
seq.append(SkeletonConv(neighbour_list, in_channels=in_channels, out_channels=in_channels,
|
||||
joint_num=joint_num, kernel_size=kernel_size - 1 if kernel_even else kernel_size,
|
||||
stride=1,
|
||||
padding=padding, padding_mode=padding_mode, bias=bias))
|
||||
seq.append(nn.PReLU() if activation == 'relu' else nn.Tanh())
|
||||
# (2T, J', D) => (2T, J', D/2)
|
||||
seq.append(SkeletonConv(neighbour_list, in_channels=in_channels, out_channels=out_channels,
|
||||
joint_num=joint_num, kernel_size=kernel_size - 1 if kernel_even else kernel_size,
|
||||
stride=1,
|
||||
padding=padding, padding_mode=padding_mode, bias=bias, add_offset=False))
|
||||
self.residual = nn.Sequential(*seq)
|
||||
|
||||
# (2T, J', D) => (2T, J', D/2)
|
||||
self.shortcut = SkeletonConv(neighbour_list, in_channels=in_channels, out_channels=out_channels,
|
||||
joint_num=joint_num, kernel_size=1, stride=1, padding=0,
|
||||
bias=True, add_offset=False)
|
||||
|
||||
if activation == 'relu':
|
||||
self.activation = nn.PReLU() if not last_layer else None
|
||||
else:
|
||||
self.activation = nn.Tanh() if not last_layer else None
|
||||
|
||||
def forward(self, input):
|
||||
output = self.common(input)
|
||||
output = self.residual(output) + self.shortcut(output)
|
||||
|
||||
if self.activation is not None:
|
||||
return self.activation(output)
|
||||
else:
|
||||
return output
|
||||
+443
@@ -0,0 +1,443 @@
|
||||
import math
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class SkeletonConv(nn.Module):
|
||||
def __init__(self, neighbour_list, in_channels, out_channels, kernel_size, joint_num, stride=1, padding=0,
|
||||
bias=True, padding_mode='zeros', add_offset=False, in_offset_channel=0):
|
||||
self.in_channels_per_joint = in_channels // joint_num
|
||||
self.out_channels_per_joint = out_channels // joint_num
|
||||
if in_channels % joint_num != 0 or out_channels % joint_num != 0:
|
||||
raise Exception('BAD')
|
||||
super(SkeletonConv, self).__init__()
|
||||
|
||||
if padding_mode == 'zeros':
|
||||
padding_mode = 'constant'
|
||||
if padding_mode == 'reflection':
|
||||
padding_mode = 'reflect'
|
||||
|
||||
self.expanded_neighbour_list = []
|
||||
self.expanded_neighbour_list_offset = []
|
||||
self.neighbour_list = neighbour_list
|
||||
self.add_offset = add_offset
|
||||
self.joint_num = joint_num
|
||||
|
||||
self.stride = stride
|
||||
self.dilation = 1
|
||||
self.groups = 1
|
||||
self.padding = padding
|
||||
self.padding_mode = padding_mode
|
||||
self._padding_repeated_twice = (padding, padding)
|
||||
|
||||
for neighbour in neighbour_list:
|
||||
expanded = []
|
||||
for k in neighbour:
|
||||
for i in range(self.in_channels_per_joint):
|
||||
expanded.append(k * self.in_channels_per_joint + i)
|
||||
self.expanded_neighbour_list.append(expanded)
|
||||
|
||||
if self.add_offset:
|
||||
self.offset_enc = SkeletonLinear(neighbour_list, in_offset_channel * len(neighbour_list), out_channels)
|
||||
|
||||
for neighbour in neighbour_list:
|
||||
expanded = []
|
||||
for k in neighbour:
|
||||
for i in range(add_offset):
|
||||
expanded.append(k * in_offset_channel + i)
|
||||
self.expanded_neighbour_list_offset.append(expanded)
|
||||
|
||||
self.weight = torch.zeros(out_channels, in_channels, kernel_size)
|
||||
if bias:
|
||||
self.bias = torch.zeros(out_channels)
|
||||
else:
|
||||
self.register_parameter('bias', None)
|
||||
|
||||
self.mask = torch.zeros_like(self.weight)
|
||||
for i, neighbour in enumerate(self.expanded_neighbour_list):
|
||||
self.mask[self.out_channels_per_joint * i: self.out_channels_per_joint * (i + 1), neighbour, ...] = 1
|
||||
self.mask = nn.Parameter(self.mask, requires_grad=False)
|
||||
|
||||
self.description = 'SkeletonConv(in_channels_per_armature={}, out_channels_per_armature={}, kernel_size={}, ' \
|
||||
'joint_num={}, stride={}, padding={}, bias={})'.format(
|
||||
in_channels // joint_num, out_channels // joint_num, kernel_size, joint_num, stride, padding, bias
|
||||
)
|
||||
|
||||
self.reset_parameters()
|
||||
|
||||
def reset_parameters(self):
|
||||
for i, neighbour in enumerate(self.expanded_neighbour_list):
|
||||
""" Use temporary variable to avoid assign to copy of slice, which might lead to unexpected result """
|
||||
tmp = torch.zeros_like(self.weight[self.out_channels_per_joint * i: self.out_channels_per_joint * (i + 1),
|
||||
neighbour, ...])
|
||||
nn.init.kaiming_uniform_(tmp, a=math.sqrt(5))
|
||||
self.weight[self.out_channels_per_joint * i: self.out_channels_per_joint * (i + 1),
|
||||
neighbour, ...] = tmp
|
||||
if self.bias is not None:
|
||||
fan_in, _ = nn.init._calculate_fan_in_and_fan_out(
|
||||
self.weight[self.out_channels_per_joint * i: self.out_channels_per_joint * (i + 1), neighbour, ...])
|
||||
bound = 1 / math.sqrt(fan_in)
|
||||
tmp = torch.zeros_like(
|
||||
self.bias[self.out_channels_per_joint * i: self.out_channels_per_joint * (i + 1)])
|
||||
nn.init.uniform_(tmp, -bound, bound)
|
||||
self.bias[self.out_channels_per_joint * i: self.out_channels_per_joint * (i + 1)] = tmp
|
||||
|
||||
self.weight = nn.Parameter(self.weight)
|
||||
if self.bias is not None:
|
||||
self.bias = nn.Parameter(self.bias)
|
||||
|
||||
def set_offset(self, offset):
|
||||
if not self.add_offset:
|
||||
raise Exception('Wrong Combination of Parameters')
|
||||
self.offset = offset.reshape(offset.shape[0], -1)
|
||||
|
||||
def forward(self, input):
|
||||
# print('SkeletonConv')
|
||||
weight_masked = self.weight * self.mask
|
||||
# print(f'input: {input.size()}')
|
||||
res = F.conv1d(F.pad(input, self._padding_repeated_twice, mode=self.padding_mode),
|
||||
weight_masked, self.bias, self.stride,
|
||||
0, self.dilation, self.groups)
|
||||
|
||||
if self.add_offset:
|
||||
offset_res = self.offset_enc(self.offset)
|
||||
offset_res = offset_res.reshape(offset_res.shape + (1, ))
|
||||
res += offset_res / 100
|
||||
# print(f'res: {res.size()}')
|
||||
return res
|
||||
|
||||
|
||||
class SkeletonLinear(nn.Module):
|
||||
def __init__(self, neighbour_list, in_channels, out_channels, extra_dim1=False):
|
||||
super(SkeletonLinear, self).__init__()
|
||||
self.neighbour_list = neighbour_list
|
||||
self.in_channels = in_channels
|
||||
self.out_channels = out_channels
|
||||
self.in_channels_per_joint = in_channels // len(neighbour_list)
|
||||
self.out_channels_per_joint = out_channels // len(neighbour_list)
|
||||
self.extra_dim1 = extra_dim1
|
||||
self.expanded_neighbour_list = []
|
||||
|
||||
for neighbour in neighbour_list:
|
||||
expanded = []
|
||||
for k in neighbour:
|
||||
for i in range(self.in_channels_per_joint):
|
||||
expanded.append(k * self.in_channels_per_joint + i)
|
||||
self.expanded_neighbour_list.append(expanded)
|
||||
|
||||
self.weight = torch.zeros(out_channels, in_channels)
|
||||
self.mask = torch.zeros(out_channels, in_channels)
|
||||
self.bias = nn.Parameter(torch.Tensor(out_channels))
|
||||
|
||||
self.reset_parameters()
|
||||
|
||||
def reset_parameters(self):
|
||||
for i, neighbour in enumerate(self.expanded_neighbour_list):
|
||||
tmp = torch.zeros_like(
|
||||
self.weight[i*self.out_channels_per_joint: (i + 1)*self.out_channels_per_joint, neighbour]
|
||||
)
|
||||
self.mask[i*self.out_channels_per_joint: (i + 1)*self.out_channels_per_joint, neighbour] = 1
|
||||
nn.init.kaiming_uniform_(tmp, a=math.sqrt(5))
|
||||
self.weight[i*self.out_channels_per_joint: (i + 1)*self.out_channels_per_joint, neighbour] = tmp
|
||||
|
||||
fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight)
|
||||
bound = 1 / math.sqrt(fan_in)
|
||||
nn.init.uniform_(self.bias, -bound, bound)
|
||||
|
||||
self.weight = nn.Parameter(self.weight)
|
||||
self.mask = nn.Parameter(self.mask, requires_grad=False)
|
||||
|
||||
def forward(self, input):
|
||||
input = input.reshape(input.shape[0], -1)
|
||||
weight_masked = self.weight * self.mask
|
||||
res = F.linear(input, weight_masked, self.bias)
|
||||
if self.extra_dim1:
|
||||
res = res.reshape(res.shape + (1,))
|
||||
return res
|
||||
|
||||
|
||||
class SkeletonPool(nn.Module):
|
||||
def __init__(self, edges, pooling_mode, channels_per_edge, last_pool=False):
|
||||
super(SkeletonPool, self).__init__()
|
||||
|
||||
if pooling_mode != 'mean':
|
||||
raise Exception('Unimplemented pooling mode in matrix_implementation')
|
||||
|
||||
self.channels_per_edge = channels_per_edge
|
||||
self.pooling_mode = pooling_mode
|
||||
self.edge_num = len(edges)
|
||||
# self.edge_num = len(edges) + 1
|
||||
self.seq_list = []
|
||||
self.pooling_list = []
|
||||
self.new_edges = []
|
||||
degree = [0] * 100 # each element represents the degree of the corresponding joint
|
||||
|
||||
for edge in edges:
|
||||
degree[edge[0]] += 1
|
||||
degree[edge[1]] += 1
|
||||
|
||||
# seq_list contains multiple sub-lists where each sub-list is an edge chain from the joint whose degree > 2 to the end effectors or joints whose degree > 2.
|
||||
def find_seq(j, seq):
|
||||
nonlocal self, degree, edges
|
||||
|
||||
if degree[j] > 2 and j != 0:
|
||||
self.seq_list.append(seq)
|
||||
seq = []
|
||||
|
||||
if degree[j] == 1:
|
||||
self.seq_list.append(seq)
|
||||
return
|
||||
|
||||
for idx, edge in enumerate(edges):
|
||||
if edge[0] == j:
|
||||
find_seq(edge[1], seq + [idx])
|
||||
|
||||
find_seq(0, [])
|
||||
# print(f'self.seq_list: {self.seq_list}')
|
||||
|
||||
for seq in self.seq_list:
|
||||
if last_pool:
|
||||
self.pooling_list.append(seq)
|
||||
continue
|
||||
if len(seq) % 2 == 1:
|
||||
self.pooling_list.append([seq[0]])
|
||||
self.new_edges.append(edges[seq[0]])
|
||||
seq = seq[1:]
|
||||
for i in range(0, len(seq), 2):
|
||||
self.pooling_list.append([seq[i], seq[i + 1]])
|
||||
self.new_edges.append([edges[seq[i]][0], edges[seq[i + 1]][1]])
|
||||
# print(f'self.pooling_list: {self.pooling_list}')
|
||||
# print(f'self.new_egdes: {self.new_edges}')
|
||||
|
||||
# add global position
|
||||
# self.pooling_list.append([self.edge_num - 1])
|
||||
|
||||
self.description = 'SkeletonPool(in_edge_num={}, out_edge_num={})'.format(
|
||||
len(edges), len(self.pooling_list)
|
||||
)
|
||||
|
||||
self.weight = torch.zeros(len(self.pooling_list) * channels_per_edge, self.edge_num * channels_per_edge)
|
||||
|
||||
for i, pair in enumerate(self.pooling_list):
|
||||
for j in pair:
|
||||
for c in range(channels_per_edge):
|
||||
self.weight[i * channels_per_edge + c, j * channels_per_edge + c] = 1.0 / len(pair)
|
||||
|
||||
self.weight = nn.Parameter(self.weight, requires_grad=False)
|
||||
|
||||
def forward(self, input: torch.Tensor):
|
||||
# print('SkeletonPool')
|
||||
# print(f'input: {input.size()}')
|
||||
# print(f'self.weight: {self.weight.size()}')
|
||||
return torch.matmul(self.weight, input)
|
||||
|
||||
|
||||
class SkeletonUnpool(nn.Module):
|
||||
def __init__(self, pooling_list, channels_per_edge):
|
||||
super(SkeletonUnpool, self).__init__()
|
||||
self.pooling_list = pooling_list
|
||||
self.input_edge_num = len(pooling_list)
|
||||
self.output_edge_num = 0
|
||||
self.channels_per_edge = channels_per_edge
|
||||
for t in self.pooling_list:
|
||||
self.output_edge_num += len(t)
|
||||
|
||||
self.description = 'SkeletonUnpool(in_edge_num={}, out_edge_num={})'.format(
|
||||
self.input_edge_num, self.output_edge_num,
|
||||
)
|
||||
|
||||
self.weight = torch.zeros(self.output_edge_num * channels_per_edge, self.input_edge_num * channels_per_edge)
|
||||
|
||||
for i, pair in enumerate(self.pooling_list):
|
||||
for j in pair:
|
||||
for c in range(channels_per_edge):
|
||||
self.weight[j * channels_per_edge + c, i * channels_per_edge + c] = 1
|
||||
|
||||
self.weight = nn.Parameter(self.weight)
|
||||
self.weight.requires_grad_(False)
|
||||
|
||||
def forward(self, input: torch.Tensor):
|
||||
# print('SkeletonUnpool')
|
||||
# print(f'input: {input.size()}')
|
||||
# print(f'self.weight: {self.weight.size()}')
|
||||
return torch.matmul(self.weight, input)
|
||||
|
||||
|
||||
"""
|
||||
Helper functions for skeleton operation
|
||||
"""
|
||||
|
||||
|
||||
def dfs(x, fa, vis, dist):
|
||||
vis[x] = 1
|
||||
for y in range(len(fa)):
|
||||
if (fa[y] == x or fa[x] == y) and vis[y] == 0:
|
||||
dist[y] = dist[x] + 1
|
||||
dfs(y, fa, vis, dist)
|
||||
|
||||
|
||||
"""
|
||||
def find_neighbor_joint(fa, threshold):
|
||||
neighbor_list = [[]]
|
||||
for x in range(1, len(fa)):
|
||||
vis = [0 for _ in range(len(fa))]
|
||||
dist = [0 for _ in range(len(fa))]
|
||||
dist[0] = 10000
|
||||
dfs(x, fa, vis, dist)
|
||||
neighbor = []
|
||||
for j in range(1, len(fa)):
|
||||
if dist[j] <= threshold:
|
||||
neighbor.append(j)
|
||||
neighbor_list.append(neighbor)
|
||||
|
||||
neighbor = [0]
|
||||
for i, x in enumerate(neighbor_list):
|
||||
if i == 0: continue
|
||||
if 1 in x:
|
||||
neighbor.append(i)
|
||||
neighbor_list[i] = [0] + neighbor_list[i]
|
||||
neighbor_list[0] = neighbor
|
||||
return neighbor_list
|
||||
|
||||
|
||||
def build_edge_topology(topology, offset):
|
||||
# get all edges (pa, child, offset)
|
||||
edges = []
|
||||
joint_num = len(topology)
|
||||
for i in range(1, joint_num):
|
||||
edges.append((topology[i], i, offset[i]))
|
||||
return edges
|
||||
"""
|
||||
|
||||
|
||||
def build_edge_topology(topology):
|
||||
# get all edges (pa, child)
|
||||
edges = []
|
||||
joint_num = len(topology)
|
||||
edges.append((0, joint_num)) # add an edge between the root joint and a virtual joint
|
||||
for i in range(1, joint_num):
|
||||
edges.append((topology[i], i))
|
||||
return edges
|
||||
|
||||
|
||||
def build_joint_topology(edges, origin_names):
|
||||
parent = []
|
||||
offset = []
|
||||
names = []
|
||||
edge2joint = []
|
||||
joint_from_edge = [] # -1 means virtual joint
|
||||
joint_cnt = 0
|
||||
out_degree = [0] * (len(edges) + 10)
|
||||
for edge in edges:
|
||||
out_degree[edge[0]] += 1
|
||||
|
||||
# add root joint
|
||||
joint_from_edge.append(-1)
|
||||
parent.append(0)
|
||||
offset.append(np.array([0, 0, 0]))
|
||||
names.append(origin_names[0])
|
||||
joint_cnt += 1
|
||||
|
||||
def make_topology(edge_idx, pa):
|
||||
nonlocal edges, parent, offset, names, edge2joint, joint_from_edge, joint_cnt
|
||||
edge = edges[edge_idx]
|
||||
if out_degree[edge[0]] > 1:
|
||||
parent.append(pa)
|
||||
offset.append(np.array([0, 0, 0]))
|
||||
names.append(origin_names[edge[1]] + '_virtual')
|
||||
edge2joint.append(-1)
|
||||
pa = joint_cnt
|
||||
joint_cnt += 1
|
||||
|
||||
parent.append(pa)
|
||||
offset.append(edge[2])
|
||||
names.append(origin_names[edge[1]])
|
||||
edge2joint.append(edge_idx)
|
||||
pa = joint_cnt
|
||||
joint_cnt += 1
|
||||
|
||||
for idx, e in enumerate(edges):
|
||||
if e[0] == edge[1]:
|
||||
make_topology(idx, pa)
|
||||
|
||||
for idx, e in enumerate(edges):
|
||||
if e[0] == 0:
|
||||
make_topology(idx, 0)
|
||||
|
||||
return parent, offset, names, edge2joint
|
||||
|
||||
|
||||
def calc_edge_mat(edges):
|
||||
edge_num = len(edges)
|
||||
# edge_mat[i][j] = distance between edge(i) and edge(j)
|
||||
edge_mat = [[100000] * edge_num for _ in range(edge_num)]
|
||||
for i in range(edge_num):
|
||||
edge_mat[i][i] = 0
|
||||
|
||||
# initialize edge_mat with direct neighbor
|
||||
for i, a in enumerate(edges):
|
||||
for j, b in enumerate(edges):
|
||||
link = 0
|
||||
for x in range(2):
|
||||
for y in range(2):
|
||||
if a[x] == b[y]:
|
||||
link = 1
|
||||
if link:
|
||||
edge_mat[i][j] = 1
|
||||
|
||||
# calculate all the pairs distance
|
||||
for k in range(edge_num):
|
||||
for i in range(edge_num):
|
||||
for j in range(edge_num):
|
||||
edge_mat[i][j] = min(edge_mat[i][j], edge_mat[i][k] + edge_mat[k][j])
|
||||
return edge_mat
|
||||
|
||||
|
||||
def find_neighbor(edges, d):
|
||||
"""
|
||||
Args:
|
||||
edges: The list contains N elements, each element represents (parent, child).
|
||||
d: Distance between edges (the distance of the same edge is 0 and the distance of adjacent edges is 1).
|
||||
|
||||
Returns:
|
||||
The list contains N elements, each element is a list of edge indices whose distance <= d.
|
||||
"""
|
||||
edge_mat = calc_edge_mat(edges)
|
||||
neighbor_list = []
|
||||
edge_num = len(edge_mat)
|
||||
for i in range(edge_num):
|
||||
neighbor = []
|
||||
for j in range(edge_num):
|
||||
if edge_mat[i][j] <= d:
|
||||
neighbor.append(j)
|
||||
neighbor_list.append(neighbor)
|
||||
|
||||
# # add neighbor for global part
|
||||
# global_part_neighbor = neighbor_list[0].copy()
|
||||
# """
|
||||
# Line #373 is buggy. Thanks @crissallan!!
|
||||
# See issue #30 (https://github.com/DeepMotionEditing/deep-motion-editing/issues/30)
|
||||
# However, fixing this bug will make it unable to load the pretrained model and
|
||||
# affect the reproducibility of quantitative error reported in the paper.
|
||||
# It is not a fatal bug so we didn't touch it and we are looking for possible solutions.
|
||||
# """
|
||||
# for i in global_part_neighbor:
|
||||
# neighbor_list[i].append(edge_num)
|
||||
# neighbor_list.append(global_part_neighbor)
|
||||
|
||||
return neighbor_list
|
||||
|
||||
|
||||
def calc_node_depth(topology):
|
||||
def dfs(node, topology):
|
||||
if topology[node] < 0:
|
||||
return 0
|
||||
return 1 + dfs(topology[node], topology)
|
||||
depth = []
|
||||
for i in range(len(topology)):
|
||||
depth.append(dfs(i, topology))
|
||||
|
||||
return depth
|
||||
+652
@@ -0,0 +1,652 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
#
|
||||
# This source code is licensed under the BSD-style license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
import numpy as np
|
||||
from typing import Optional, Union
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
Device = Union[str, torch.device]
|
||||
|
||||
|
||||
"""
|
||||
The transformation matrices returned from the functions in this file assume
|
||||
the points on which the transformation will be applied are column vectors.
|
||||
i.e. the R matrix is structured as
|
||||
|
||||
R = [
|
||||
[Rxx, Rxy, Rxz],
|
||||
[Ryx, Ryy, Ryz],
|
||||
[Rzx, Rzy, Rzz],
|
||||
] # (3, 3)
|
||||
|
||||
This matrix can be applied to column vectors by post multiplication
|
||||
by the points e.g.
|
||||
|
||||
points = [[0], [1], [2]] # (3 x 1) xyz coordinates of a point
|
||||
transformed_points = R * points
|
||||
|
||||
To apply the same matrix to points which are row vectors, the R matrix
|
||||
can be transposed and pre multiplied by the points:
|
||||
|
||||
e.g.
|
||||
points = [[0, 1, 2]] # (1 x 3) xyz coordinates of a point
|
||||
transformed_points = points * R.transpose(1, 0)
|
||||
"""
|
||||
|
||||
|
||||
def quaternion_to_matrix(quaternions: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Convert rotations given as quaternions to rotation matrices.
|
||||
|
||||
Args:
|
||||
quaternions: quaternions with real part first,
|
||||
as tensor of shape (..., 4).
|
||||
|
||||
Returns:
|
||||
Rotation matrices as tensor of shape (..., 3, 3).
|
||||
"""
|
||||
r, i, j, k = torch.unbind(quaternions, -1)
|
||||
two_s = 2.0 / (quaternions * quaternions).sum(-1)
|
||||
|
||||
o = torch.stack(
|
||||
(
|
||||
1 - two_s * (j * j + k * k),
|
||||
two_s * (i * j - k * r),
|
||||
two_s * (i * k + j * r),
|
||||
two_s * (i * j + k * r),
|
||||
1 - two_s * (i * i + k * k),
|
||||
two_s * (j * k - i * r),
|
||||
two_s * (i * k - j * r),
|
||||
two_s * (j * k + i * r),
|
||||
1 - two_s * (i * i + j * j),
|
||||
),
|
||||
-1,
|
||||
)
|
||||
return o.reshape(quaternions.shape[:-1] + (3, 3))
|
||||
|
||||
|
||||
def _copysign(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Return a tensor where each element has the absolute value taken from the,
|
||||
corresponding element of a, with sign taken from the corresponding
|
||||
element of b. This is like the standard copysign floating-point operation,
|
||||
but is not careful about negative 0 and NaN.
|
||||
|
||||
Args:
|
||||
a: source tensor.
|
||||
b: tensor whose signs will be used, of the same shape as a.
|
||||
|
||||
Returns:
|
||||
Tensor of the same shape as a with the signs of b.
|
||||
"""
|
||||
signs_differ = (a < 0) != (b < 0)
|
||||
return torch.where(signs_differ, -a, a)
|
||||
|
||||
|
||||
def _sqrt_positive_part(x: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Returns torch.sqrt(torch.max(0, x))
|
||||
but with a zero subgradient where x is 0.
|
||||
"""
|
||||
ret = torch.zeros_like(x)
|
||||
positive_mask = x > 0
|
||||
ret[positive_mask] = torch.sqrt(x[positive_mask])
|
||||
return ret
|
||||
|
||||
|
||||
def matrix_to_quaternion(matrix: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Convert rotations given as rotation matrices to quaternions.
|
||||
|
||||
Args:
|
||||
matrix: Rotation matrices as tensor of shape (..., 3, 3).
|
||||
|
||||
Returns:
|
||||
quaternions with real part first, as tensor of shape (..., 4).
|
||||
"""
|
||||
if matrix.size(-1) != 3 or matrix.size(-2) != 3:
|
||||
raise ValueError(f"Invalid rotation matrix shape {matrix.shape}.")
|
||||
|
||||
batch_dim = matrix.shape[:-2]
|
||||
m00, m01, m02, m10, m11, m12, m20, m21, m22 = torch.unbind(
|
||||
matrix.reshape(batch_dim + (9,)), dim=-1
|
||||
)
|
||||
|
||||
q_abs = _sqrt_positive_part(
|
||||
torch.stack(
|
||||
[
|
||||
1.0 + m00 + m11 + m22,
|
||||
1.0 + m00 - m11 - m22,
|
||||
1.0 - m00 + m11 - m22,
|
||||
1.0 - m00 - m11 + m22,
|
||||
],
|
||||
dim=-1,
|
||||
)
|
||||
)
|
||||
|
||||
# we produce the desired quaternion multiplied by each of r, i, j, k
|
||||
quat_by_rijk = torch.stack(
|
||||
[
|
||||
torch.stack([q_abs[..., 0] ** 2, m21 - m12, m02 - m20, m10 - m01], dim=-1),
|
||||
torch.stack([m21 - m12, q_abs[..., 1] ** 2, m10 + m01, m02 + m20], dim=-1),
|
||||
torch.stack([m02 - m20, m10 + m01, q_abs[..., 2] ** 2, m12 + m21], dim=-1),
|
||||
torch.stack([m10 - m01, m20 + m02, m21 + m12, q_abs[..., 3] ** 2], dim=-1),
|
||||
],
|
||||
dim=-2,
|
||||
)
|
||||
|
||||
# We floor here at 0.1 but the exact level is not important; if q_abs is small,
|
||||
# the candidate won't be picked.
|
||||
flr = torch.tensor(0.1).to(dtype=q_abs.dtype, device=q_abs.device)
|
||||
quat_candidates = quat_by_rijk / (2.0 * q_abs[..., None].max(flr))
|
||||
|
||||
# if not for numerical problems, quat_candidates[i] should be same (up to a sign),
|
||||
# forall i; we pick the best-conditioned one (with the largest denominator)
|
||||
|
||||
return quat_candidates[
|
||||
F.one_hot(q_abs.argmax(dim=-1), num_classes=4) > 0.5, : # pyre-ignore[16]
|
||||
].reshape(batch_dim + (4,))
|
||||
|
||||
|
||||
def _axis_angle_rotation(axis: str, angle: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Return the rotation matrices for one of the rotations about an axis
|
||||
of which Euler angles describe, for each value of the angle given.
|
||||
|
||||
Args:
|
||||
axis: Axis label "X" or "Y or "Z".
|
||||
angle: any shape tensor of Euler angles in radians
|
||||
|
||||
Returns:
|
||||
Rotation matrices as tensor of shape (..., 3, 3).
|
||||
"""
|
||||
|
||||
cos = torch.cos(angle)
|
||||
sin = torch.sin(angle)
|
||||
one = torch.ones_like(angle)
|
||||
zero = torch.zeros_like(angle)
|
||||
|
||||
if axis == "X":
|
||||
R_flat = (one, zero, zero, zero, cos, -sin, zero, sin, cos)
|
||||
elif axis == "Y":
|
||||
R_flat = (cos, zero, sin, zero, one, zero, -sin, zero, cos)
|
||||
elif axis == "Z":
|
||||
R_flat = (cos, -sin, zero, sin, cos, zero, zero, zero, one)
|
||||
else:
|
||||
raise ValueError("letter must be either X, Y or Z.")
|
||||
|
||||
return torch.stack(R_flat, -1).reshape(angle.shape + (3, 3))
|
||||
|
||||
|
||||
def euler_angles_to_matrix(euler_angles: torch.Tensor, convention: str) -> torch.Tensor:
|
||||
"""
|
||||
Convert rotations given as Euler angles in radians to rotation matrices.
|
||||
|
||||
Args:
|
||||
euler_angles: Euler angles in radians as tensor of shape (..., 3).
|
||||
convention: Convention string of three uppercase letters from
|
||||
{"X", "Y", and "Z"}.
|
||||
|
||||
Returns:
|
||||
Rotation matrices as tensor of shape (..., 3, 3).
|
||||
"""
|
||||
if euler_angles.dim() == 0 or euler_angles.shape[-1] != 3:
|
||||
raise ValueError("Invalid input euler angles.")
|
||||
if len(convention) != 3:
|
||||
raise ValueError("Convention must have 3 letters.")
|
||||
if convention[1] in (convention[0], convention[2]):
|
||||
raise ValueError(f"Invalid convention {convention}.")
|
||||
for letter in convention:
|
||||
if letter not in ("X", "Y", "Z"):
|
||||
raise ValueError(f"Invalid letter {letter} in convention string.")
|
||||
matrices = [
|
||||
_axis_angle_rotation(c, e)
|
||||
for c, e in zip(convention, torch.unbind(euler_angles, -1))
|
||||
]
|
||||
# return functools.reduce(torch.matmul, matrices)
|
||||
return torch.matmul(torch.matmul(matrices[0], matrices[1]), matrices[2])
|
||||
|
||||
|
||||
def _angle_from_tan(
|
||||
axis: str, other_axis: str, data, horizontal: bool, tait_bryan: bool
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Extract the first or third Euler angle from the two members of
|
||||
the matrix which are positive constant times its sine and cosine.
|
||||
|
||||
Args:
|
||||
axis: Axis label "X" or "Y or "Z" for the angle we are finding.
|
||||
other_axis: Axis label "X" or "Y or "Z" for the middle axis in the
|
||||
convention.
|
||||
data: Rotation matrices as tensor of shape (..., 3, 3).
|
||||
horizontal: Whether we are looking for the angle for the third axis,
|
||||
which means the relevant entries are in the same row of the
|
||||
rotation matrix. If not, they are in the same column.
|
||||
tait_bryan: Whether the first and third axes in the convention differ.
|
||||
|
||||
Returns:
|
||||
Euler Angles in radians for each matrix in data as a tensor
|
||||
of shape (...).
|
||||
"""
|
||||
|
||||
i1, i2 = {"X": (2, 1), "Y": (0, 2), "Z": (1, 0)}[axis]
|
||||
if horizontal:
|
||||
i2, i1 = i1, i2
|
||||
even = (axis + other_axis) in ["XY", "YZ", "ZX"]
|
||||
if horizontal == even:
|
||||
return torch.atan2(data[..., i1], data[..., i2])
|
||||
if tait_bryan:
|
||||
return torch.atan2(-data[..., i2], data[..., i1])
|
||||
return torch.atan2(data[..., i2], -data[..., i1])
|
||||
|
||||
|
||||
def _index_from_letter(letter: str) -> int:
|
||||
if letter == "X":
|
||||
return 0
|
||||
if letter == "Y":
|
||||
return 1
|
||||
if letter == "Z":
|
||||
return 2
|
||||
raise ValueError("letter must be either X, Y or Z.")
|
||||
|
||||
|
||||
def matrix_to_euler_angles(matrix: torch.Tensor, convention: str) -> torch.Tensor:
|
||||
"""
|
||||
Convert rotations given as rotation matrices to Euler angles in radians.
|
||||
|
||||
Args:
|
||||
matrix: Rotation matrices as tensor of shape (..., 3, 3).
|
||||
convention: Convention string of three uppercase letters.
|
||||
|
||||
Returns:
|
||||
Euler angles in radians as tensor of shape (..., 3).
|
||||
"""
|
||||
if len(convention) != 3:
|
||||
raise ValueError("Convention must have 3 letters.")
|
||||
if convention[1] in (convention[0], convention[2]):
|
||||
raise ValueError(f"Invalid convention {convention}.")
|
||||
for letter in convention:
|
||||
if letter not in ("X", "Y", "Z"):
|
||||
raise ValueError(f"Invalid letter {letter} in convention string.")
|
||||
if matrix.size(-1) != 3 or matrix.size(-2) != 3:
|
||||
raise ValueError(f"Invalid rotation matrix shape {matrix.shape}.")
|
||||
i0 = _index_from_letter(convention[0])
|
||||
i2 = _index_from_letter(convention[2])
|
||||
tait_bryan = i0 != i2
|
||||
if tait_bryan:
|
||||
central_angle = torch.asin(
|
||||
matrix[..., i0, i2] * (-1.0 if i0 - i2 in [-1, 2] else 1.0)
|
||||
)
|
||||
else:
|
||||
central_angle = torch.acos(matrix[..., i0, i0])
|
||||
|
||||
o = (
|
||||
_angle_from_tan(
|
||||
convention[0], convention[1], matrix[..., i2], False, tait_bryan
|
||||
),
|
||||
central_angle,
|
||||
_angle_from_tan(
|
||||
convention[2], convention[1], matrix[..., i0, :], True, tait_bryan
|
||||
),
|
||||
)
|
||||
return torch.stack(o, -1)
|
||||
|
||||
|
||||
def random_quaternions(
|
||||
n: int, dtype: Optional[torch.dtype] = None, device: Optional[Device] = None
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Generate random quaternions representing rotations,
|
||||
i.e. versors with nonnegative real part.
|
||||
|
||||
Args:
|
||||
n: Number of quaternions in a batch to return.
|
||||
dtype: Type to return.
|
||||
device: Desired device of returned tensor. Default:
|
||||
uses the current device for the default tensor type.
|
||||
|
||||
Returns:
|
||||
Quaternions as tensor of shape (N, 4).
|
||||
"""
|
||||
if isinstance(device, str):
|
||||
device = torch.device(device)
|
||||
o = torch.randn((n, 4), dtype=dtype, device=device)
|
||||
s = (o * o).sum(1)
|
||||
o = o / _copysign(torch.sqrt(s), o[:, 0])[:, None]
|
||||
return o
|
||||
|
||||
|
||||
def random_rotations(
|
||||
n: int, dtype: Optional[torch.dtype] = None, device: Optional[Device] = None
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Generate random rotations as 3x3 rotation matrices.
|
||||
|
||||
Args:
|
||||
n: Number of rotation matrices in a batch to return.
|
||||
dtype: Type to return.
|
||||
device: Device of returned tensor. Default: if None,
|
||||
uses the current device for the default tensor type.
|
||||
|
||||
Returns:
|
||||
Rotation matrices as tensor of shape (n, 3, 3).
|
||||
"""
|
||||
quaternions = random_quaternions(n, dtype=dtype, device=device)
|
||||
return quaternion_to_matrix(quaternions)
|
||||
|
||||
|
||||
def random_rotation(
|
||||
dtype: Optional[torch.dtype] = None, device: Optional[Device] = None
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Generate a single random 3x3 rotation matrix.
|
||||
|
||||
Args:
|
||||
dtype: Type to return
|
||||
device: Device of returned tensor. Default: if None,
|
||||
uses the current device for the default tensor type
|
||||
|
||||
Returns:
|
||||
Rotation matrix as tensor of shape (3, 3).
|
||||
"""
|
||||
return random_rotations(1, dtype, device)[0]
|
||||
|
||||
|
||||
def standardize_quaternion(quaternions: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Convert a unit quaternion to a standard form: one in which the real
|
||||
part is non negative.
|
||||
|
||||
Args:
|
||||
quaternions: Quaternions with real part first,
|
||||
as tensor of shape (..., 4).
|
||||
|
||||
Returns:
|
||||
Standardized quaternions as tensor of shape (..., 4).
|
||||
"""
|
||||
return torch.where(quaternions[..., 0:1] < 0, -quaternions, quaternions)
|
||||
|
||||
|
||||
def quaternion_raw_multiply(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Multiply two quaternions.
|
||||
Usual torch rules for broadcasting apply.
|
||||
|
||||
Args:
|
||||
a: Quaternions as tensor of shape (..., 4), real part first.
|
||||
b: Quaternions as tensor of shape (..., 4), real part first.
|
||||
|
||||
Returns:
|
||||
The product of a and b, a tensor of quaternions shape (..., 4).
|
||||
"""
|
||||
aw, ax, ay, az = torch.unbind(a, -1)
|
||||
bw, bx, by, bz = torch.unbind(b, -1)
|
||||
ow = aw * bw - ax * bx - ay * by - az * bz
|
||||
ox = aw * bx + ax * bw + ay * bz - az * by
|
||||
oy = aw * by - ax * bz + ay * bw + az * bx
|
||||
oz = aw * bz + ax * by - ay * bx + az * bw
|
||||
return torch.stack((ow, ox, oy, oz), -1)
|
||||
|
||||
|
||||
def quaternion_multiply(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Multiply two quaternions representing rotations, returning the quaternion
|
||||
representing their composition, i.e. the versor with nonnegative real part.
|
||||
Usual torch rules for broadcasting apply.
|
||||
|
||||
Args:
|
||||
a: Quaternions as tensor of shape (..., 4), real part first.
|
||||
b: Quaternions as tensor of shape (..., 4), real part first.
|
||||
|
||||
Returns:
|
||||
The product of a and b, a tensor of quaternions of shape (..., 4).
|
||||
"""
|
||||
ab = quaternion_raw_multiply(a, b)
|
||||
return standardize_quaternion(ab)
|
||||
|
||||
|
||||
def quaternion_invert(quaternion: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Given a quaternion representing rotation, get the quaternion representing
|
||||
its inverse.
|
||||
|
||||
Args:
|
||||
quaternion: Quaternions as tensor of shape (..., 4), with real part
|
||||
first, which must be versors (unit quaternions).
|
||||
|
||||
Returns:
|
||||
The inverse, a tensor of quaternions of shape (..., 4).
|
||||
"""
|
||||
|
||||
scaling = torch.tensor([1, -1, -1, -1], device=quaternion.device)
|
||||
return quaternion * scaling
|
||||
|
||||
|
||||
def quaternion_apply(quaternion: torch.Tensor, point: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Apply the rotation given by a quaternion to a 3D point.
|
||||
Usual torch rules for broadcasting apply.
|
||||
|
||||
Args:
|
||||
quaternion: Tensor of quaternions, real part first, of shape (..., 4).
|
||||
point: Tensor of 3D points of shape (..., 3).
|
||||
|
||||
Returns:
|
||||
Tensor of rotated points of shape (..., 3).
|
||||
"""
|
||||
if point.size(-1) != 3:
|
||||
raise ValueError(f"Points are not in 3D, {point.shape}.")
|
||||
real_parts = point.new_zeros(point.shape[:-1] + (1,))
|
||||
point_as_quaternion = torch.cat((real_parts, point), -1)
|
||||
out = quaternion_raw_multiply(
|
||||
quaternion_raw_multiply(quaternion, point_as_quaternion),
|
||||
quaternion_invert(quaternion),
|
||||
)
|
||||
return out[..., 1:]
|
||||
|
||||
|
||||
def axis_angle_to_matrix(axis_angle: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Convert rotations given as axis/angle to rotation matrices.
|
||||
|
||||
Args:
|
||||
axis_angle: Rotations given as a vector in axis angle form,
|
||||
as a tensor of shape (..., 3), where the magnitude is
|
||||
the angle turned anticlockwise in radians around the
|
||||
vector's direction.
|
||||
|
||||
Returns:
|
||||
Rotation matrices as tensor of shape (..., 3, 3).
|
||||
"""
|
||||
return quaternion_to_matrix(axis_angle_to_quaternion(axis_angle))
|
||||
|
||||
|
||||
def matrix_to_axis_angle(matrix: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Convert rotations given as rotation matrices to axis/angle.
|
||||
|
||||
Args:
|
||||
matrix: Rotation matrices as tensor of shape (..., 3, 3).
|
||||
|
||||
Returns:
|
||||
Rotations given as a vector in axis angle form, as a tensor
|
||||
of shape (..., 3), where the magnitude is the angle
|
||||
turned anticlockwise in radians around the vector's
|
||||
direction.
|
||||
"""
|
||||
return quaternion_to_axis_angle(matrix_to_quaternion(matrix))
|
||||
|
||||
|
||||
def axis_angle_to_quaternion(axis_angle: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Convert rotations given as axis/angle to quaternions.
|
||||
|
||||
Args:
|
||||
axis_angle: Rotations given as a vector in axis angle form,
|
||||
as a tensor of shape (..., 3), where the magnitude is
|
||||
the angle turned anticlockwise in radians around the
|
||||
vector's direction.
|
||||
|
||||
Returns:
|
||||
quaternions with real part first, as tensor of shape (..., 4).
|
||||
"""
|
||||
angles = torch.norm(axis_angle, p=2, dim=-1, keepdim=True)
|
||||
|
||||
half_angles = angles * 0.5
|
||||
eps = 1e-6
|
||||
small_angles = angles.abs() < eps
|
||||
sin_half_angles_over_angles = torch.empty_like(angles)
|
||||
sin_half_angles_over_angles[~small_angles] = (
|
||||
torch.sin(half_angles[~small_angles]) / angles[~small_angles]
|
||||
)
|
||||
# for x small, sin(x/2) is about x/2 - (x/2)^3/6
|
||||
# so sin(x/2)/x is about 1/2 - (x*x)/48
|
||||
sin_half_angles_over_angles[small_angles] = (
|
||||
0.5 - (angles[small_angles] * angles[small_angles]) / 48
|
||||
)
|
||||
quaternions = torch.cat(
|
||||
[torch.cos(half_angles), axis_angle * sin_half_angles_over_angles], dim=-1
|
||||
)
|
||||
return quaternions
|
||||
|
||||
|
||||
def quat_to_aa(quaternions: torch.Tensor) -> torch.Tensor:
|
||||
|
||||
norms = torch.norm(quaternions[..., 1:], p=2, dim=-1, keepdim=True)
|
||||
half_angles = torch.atan2(norms, quaternions[..., :1])
|
||||
angles = 2 * half_angles
|
||||
|
||||
# if angles.max() > torch.pi:
|
||||
# angles = angles % torch.pi
|
||||
# half_angles = angles / 2
|
||||
|
||||
eps = 1e-6
|
||||
small_angles = angles.abs() < eps
|
||||
sin_half_angles_over_angles = torch.empty_like(angles)
|
||||
|
||||
# sin_half_angles_over_angles[~small_angles] = (
|
||||
# torch.sin(half_angles[~small_angles]) / angles[~small_angles])
|
||||
|
||||
sin_half_angles_over_angles[~small_angles] = (
|
||||
torch.tensor(np.sin(half_angles[~small_angles].cpu()).to(angles) / angles[~small_angles]))
|
||||
|
||||
# for x small, sin(x/2) is about x/2 - (x/2)^3/6
|
||||
# so sin(x/2)/x is about 1/2 - (x*x)/48
|
||||
sin_half_angles_over_angles[small_angles] = (
|
||||
0.5 - (angles[small_angles] * angles[small_angles]) / 48
|
||||
)
|
||||
return quaternions[..., 1:] / sin_half_angles_over_angles
|
||||
|
||||
|
||||
def quaternion_to_axis_angle(quaternions: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Convert rotations given as quaternions to axis/angle.
|
||||
|
||||
Args:
|
||||
quaternions: quaternions with real part first,
|
||||
as tensor of shape (..., 4).
|
||||
|
||||
Returns:
|
||||
Rotations given as a vector in axis angle form, as a tensor
|
||||
of shape (..., 3), where the magnitude is the angle
|
||||
turned anticlockwise in radians around the vector's
|
||||
direction.
|
||||
"""
|
||||
norms = torch.norm(quaternions[..., 1:], p=2, dim=-1, keepdim=True)
|
||||
half_angles = torch.atan2(norms, quaternions[..., :1])
|
||||
angles = 2 * half_angles
|
||||
|
||||
eps = 1e-6
|
||||
small_angles = angles.abs() < eps
|
||||
sin_half_angles_over_angles = torch.empty_like(angles)
|
||||
|
||||
sin_half_angles_over_angles[~small_angles] = (
|
||||
torch.sin(half_angles[~small_angles]) / angles[~small_angles])
|
||||
|
||||
# for x small, sin(x/2) is about x/2 - (x/2)^3/6
|
||||
# so sin(x/2)/x is about 1/2 - (x*x)/48
|
||||
sin_half_angles_over_angles[small_angles] = (
|
||||
0.5 - (angles[small_angles] * angles[small_angles]) / 48
|
||||
)
|
||||
return quaternions[..., 1:] / sin_half_angles_over_angles
|
||||
|
||||
|
||||
def rotation_6d_to_matrix(d6: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Converts 6D rotation representation by Zhou et al. [1] to rotation matrix
|
||||
using Gram--Schmidt orthogonalization per Section B of [1].
|
||||
Args:
|
||||
d6: 6D rotation representation, of size (*, 6)
|
||||
|
||||
Returns:
|
||||
batch of rotation matrices of size (*, 3, 3)
|
||||
|
||||
[1] Zhou, Y., Barnes, C., Lu, J., Yang, J., & Li, H.
|
||||
On the Continuity of Rotation Representations in Neural Networks.
|
||||
IEEE Conference on Computer Vision and Pattern Recognition, 2019.
|
||||
Retrieved from http://arxiv.org/abs/1812.07035
|
||||
"""
|
||||
|
||||
a1, a2 = d6[..., :3], d6[..., 3:]
|
||||
b1 = F.normalize(a1, dim=-1)
|
||||
b2 = a2 - (b1 * a2).sum(-1, keepdim=True) * b1
|
||||
b2 = F.normalize(b2, dim=-1)
|
||||
b3 = torch.cross(b1, b2, dim=-1)
|
||||
return torch.stack((b1, b2, b3), dim=-2)
|
||||
|
||||
|
||||
def matrix_to_rotation_6d(matrix: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Converts rotation matrices to 6D rotation representation by Zhou et al. [1]
|
||||
by dropping the last row. Note that 6D representation is not unique.
|
||||
Args:
|
||||
matrix: batch of rotation matrices of size (*, 3, 3)
|
||||
|
||||
Returns:
|
||||
6D rotation representation, of size (*, 6)
|
||||
|
||||
[1] Zhou, Y., Barnes, C., Lu, J., Yang, J., & Li, H.
|
||||
On the Continuity of Rotation Representations in Neural Networks.
|
||||
IEEE Conference on Computer Vision and Pattern Recognition, 2019.
|
||||
Retrieved from http://arxiv.org/abs/1812.07035
|
||||
"""
|
||||
batch_dim = matrix.size()[:-2]
|
||||
return matrix[..., :2, :].clone().reshape(batch_dim + (6,))
|
||||
|
||||
|
||||
def batch_rodrigues(rot_vecs, epsilon=1e-8, dtype=torch.float32):
|
||||
''' 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 = rot_vecs.device
|
||||
|
||||
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
|
||||
+473
@@ -0,0 +1,473 @@
|
||||
import os
|
||||
import cv2
|
||||
import numpy as np
|
||||
import plyfile
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from human_body_prior.tools.omni_tools import copy2cpu as c2c
|
||||
from scipy.spatial.transform import Rotation as R
|
||||
from scipy.spatial.transform import Slerp
|
||||
|
||||
import holden.BVH as BVH
|
||||
from holden.Animation import Animation
|
||||
from holden.Quaternions import Quaternions
|
||||
from typing import List, Tuple, NewType
|
||||
import math
|
||||
Tensor = NewType('Tensor', torch.Tensor)
|
||||
|
||||
SMPL_JOINTS = {
|
||||
'Pelvis': 0,
|
||||
'L_Hip': 1, 'L_Knee': 4, 'L_Ankle': 7, 'L_Foot': 10,
|
||||
'R_Hip': 2, 'R_Knee': 5, 'R_Ankle': 8, 'R_Foot': 11,
|
||||
'Spine1': 3, 'Spine2': 6, 'Spine3': 9, 'Neck': 12, 'Head': 15,
|
||||
'L_Collar': 13, 'L_Shoulder': 16, 'L_Elbow': 18, 'L_Wrist': 20, 'L_Hand': 22,
|
||||
'R_Collar': 14, 'R_Shoulder': 17, 'R_Elbow': 19, 'R_Wrist': 21, 'R_Hand': 23
|
||||
}
|
||||
|
||||
MANO_RIGHT_JOINTS = {'R_Wrist':0, 'rindex0':1, 'rindex1':2, 'rindex2':3, 'rmiddle0':4, 'rmiddle1':5, 'rmiddle2':6,
|
||||
'rpinky0':7, 'rpinky1':8, 'rpinky2':9, 'rring0':10, 'rring1':11, 'rring2':12, 'rthumb0':13, 'rthumb1':14, 'rthumb2':15}
|
||||
|
||||
|
||||
FOOT_IDX = [SMPL_JOINTS['L_Ankle'], SMPL_JOINTS['R_Ankle'], SMPL_JOINTS['L_Foot'], SMPL_JOINTS['R_Foot']]
|
||||
|
||||
CONTACTS_IDX = [SMPL_JOINTS['L_Ankle'], SMPL_JOINTS['R_Ankle'],
|
||||
SMPL_JOINTS['L_Foot'], SMPL_JOINTS['R_Foot'],
|
||||
SMPL_JOINTS['L_Wrist'], SMPL_JOINTS['R_Wrist'],
|
||||
SMPL_JOINTS['L_Knee'], SMPL_JOINTS['R_Knee']]
|
||||
|
||||
|
||||
def align_joints(tensor, smpl_to_bvh=True):
|
||||
"""
|
||||
Args:
|
||||
tensor (T x J x D): The 3D torch tensor we need to process.
|
||||
|
||||
Returns:
|
||||
The 3D tensor whose joint order is compatible for bvh export.
|
||||
"""
|
||||
order = list(MANO_RIGHT_JOINTS.values())
|
||||
result = tensor.clone() if isinstance(tensor, torch.Tensor) else tensor.copy()
|
||||
for i in range(len(order)):
|
||||
if smpl_to_bvh:
|
||||
result[:, i] = tensor[:, order[i]]
|
||||
else:
|
||||
result[:, order[i]] = tensor[:, i]
|
||||
|
||||
return result
|
||||
|
||||
def ho3d_cam_extrinsics(cam_id):
|
||||
|
||||
# GPMF case
|
||||
extrinsic_0 = np.array([[-6.409640608401522277e-01, -2.650417602434713693e-01, 7.203595893984829912e-01, -5.675702013852904626e-01],
|
||||
[-7.601847604882365772e-01, 3.490763992765258950e-01, -5.479642300274587541e-01, 3.191066048126701693e-01],
|
||||
[-1.062271275746045074e-01, -8.988317600058360890e-01, -4.252261334538034454e-01, 3.074042501856224519e-01],
|
||||
[0, 0, 0, 1]])
|
||||
|
||||
extrinsic_1 = np.array([[-4.228051957326409149e-01, -2.078562115810418665e-01, -8.820609739517225600e-01, 4.637252480950287414e-01],
|
||||
[9.057441941962964815e-01, -6.536944939351106709e-02, -4.187532564239835331e-01, 1.535094236870509776e-01],
|
||||
[2.938062526878559150e-02, -9.759726586299259932e-01, 2.159030997129234852e-01, -1.448621926159473529e-02],
|
||||
[0, 0, 0, 1]])
|
||||
|
||||
extrinsic_2 = np.array([[9.846039213996782280e-01, 6.828612271512826681e-02, -1.609102961556882483e-01, 9.897043274267488394e-03],
|
||||
[1.747949086333730750e-01, -3.774480817251544273e-01, 9.093842342584639304e-01, -4.349111883283850455e-01],
|
||||
[1.363040801985337556e-03, -9.235095836246293155e-01, -3.835729279722913110e-01, 2.864983800867586528e-01],
|
||||
[0, 0, 0, 1]])
|
||||
|
||||
extrinsic_3 = np.array([[4.290309538248047727e-01, 8.850724776239779490e-02, 8.989432172021233347e-01, -4.754434780295317409e-01],
|
||||
[-9.018737295370956586e-01, 9.767471789067121157e-02, 4.208128152229064223e-01, -2.166588733650218201e-01],
|
||||
[-5.055904104142360661e-02, -9.912749954369395322e-01, 1.217278390117415843e-01, 6.127903398258782025e-02],
|
||||
[0, 0, 0, 1]])
|
||||
|
||||
extrinsic_4 = np.array([[-9.650817552456496529e-01, 5.435035576465601509e-02, 2.562484039369455346e-01, -9.846218864128720993e-02],
|
||||
[-2.604576318061128104e-01, -3.033325236743535935e-01, -9.165976227960112022e-01, 3.948193893157375123e-01],
|
||||
[2.791106816171780650e-02, -9.513334951054777111e-01, 3.068966493210200097e-01, -1.383363501223186694e-02],
|
||||
[0, 0, 0, 1]])
|
||||
|
||||
R, t = eval(f"extrinsic_{cam_id}")[:3, :3], eval(f"extrinsic_{cam_id}")[:3, 3]
|
||||
|
||||
return R, t
|
||||
|
||||
|
||||
def slerp(quat, trans, key_times, times, mask=True):
|
||||
"""
|
||||
Args:
|
||||
quat: (T x J x 4)
|
||||
trans: (T x 3)
|
||||
"""
|
||||
if mask:
|
||||
quat = c2c(quat[key_times])
|
||||
trans = c2c(trans[key_times])
|
||||
else:
|
||||
quat = c2c(quat)
|
||||
trans = c2c(trans)
|
||||
|
||||
quats = []
|
||||
for j in range(quat.shape[1]):
|
||||
key_rots = R.from_quat(quat[:, j])
|
||||
s = Slerp(key_times, key_rots)
|
||||
interp_rots = s(times)
|
||||
quats.append(interp_rots.as_quat())
|
||||
slerp_quat = np.stack(quats, axis=1)
|
||||
|
||||
lerp_trans = np.zeros((len(times), 3))
|
||||
for i in range(3):
|
||||
lerp_trans[:, i] = np.interp(times, key_times, trans[:, i])
|
||||
|
||||
return slerp_quat, lerp_trans
|
||||
|
||||
|
||||
def compute_orient_angle(matrix, traj):
|
||||
"""
|
||||
Args:
|
||||
matrix (N x T x 3 x 3): The 3D rotation matrix at the root joint
|
||||
traj (N x T x 3): The trajectory to align
|
||||
"""
|
||||
forward = matrix[:, :, :, 2].clone()
|
||||
forward[:, :, 2] = 0
|
||||
forward = F.normalize(forward, dim=-1) # normalized forward vector (N, T, 3)
|
||||
|
||||
traj[:, :, 2] = 0 # make sure the trajectory is projected to the plane
|
||||
|
||||
# first steps is forward diff
|
||||
init_tan = traj[:, 1:2] - traj[:, :1]
|
||||
# middle steps are second order
|
||||
middle_tan = (traj[:, 2:] - traj[:, 0:-2]) / 2
|
||||
# last step is backward diff
|
||||
final_tan = traj[:, -1:] - traj[:, -2:-1]
|
||||
|
||||
tangent = torch.cat([init_tan, middle_tan, final_tan], dim=1)
|
||||
tangent = F.normalize(tangent, dim=-1) # normalized tangent vector (N, T, 3)
|
||||
|
||||
cos = torch.sum(forward * tangent, dim=-1)
|
||||
|
||||
return cos
|
||||
|
||||
|
||||
def compute_trajectory(velocity, up, origin, dt, up_axis='z'):
|
||||
"""
|
||||
Args:
|
||||
velocity: (B, T, 3)
|
||||
up: (B, T)
|
||||
origin: (B, 3)
|
||||
up_axis: x, y, or z
|
||||
|
||||
Returns:
|
||||
trajectory: (B, T, 3)
|
||||
"""
|
||||
ordermap = {
|
||||
'x': 0,
|
||||
'y': 1,
|
||||
'z': 2,
|
||||
}
|
||||
v_axis = [x for x in ordermap.values() if x != ordermap[up_axis]]
|
||||
|
||||
origin = origin.unsqueeze(1) # (B, 3) => (B, 1, 3)
|
||||
trajectory = origin.repeat(1, up.shape[1], 1) # (B, 1, 3) => (B, T, 3)
|
||||
|
||||
for t in range(1, up.shape[1]):
|
||||
trajectory[:, t, v_axis[0]] = trajectory[:, t - 1, v_axis[0]] + velocity[:, t - 1, v_axis[0]] * dt
|
||||
trajectory[:, t, v_axis[1]] = trajectory[:, t - 1, v_axis[1]] + velocity[:, t - 1, v_axis[1]] * dt
|
||||
|
||||
trajectory[:, :, ordermap[up_axis]] = up
|
||||
|
||||
return trajectory
|
||||
|
||||
|
||||
def build_canonical_frame(forward, up_axis='z'):
|
||||
"""
|
||||
Args:
|
||||
forward: (..., 3)
|
||||
|
||||
Returns:
|
||||
frame: (..., 3, 3)
|
||||
"""
|
||||
forward[..., 'xyz'.index(up_axis)] = 0
|
||||
forward = F.normalize(forward, dim=-1) # normalized forward vector
|
||||
|
||||
up = torch.zeros_like(forward)
|
||||
up[..., 'xyz'.index(up_axis)] = 1 # normalized up vector
|
||||
right = torch.cross(up, forward)
|
||||
frame = torch.stack((right, up, forward), dim=-1) # canonical frame
|
||||
|
||||
return frame
|
||||
|
||||
|
||||
def estimate_linear_velocity(data_seq, dt):
|
||||
'''
|
||||
Given some batched data sequences of T timesteps in the shape (B, T, ...), estimates
|
||||
the velocity for the middle T-2 steps using a second order central difference scheme.
|
||||
The first and last frames are with forward and backward first-order
|
||||
differences, respectively
|
||||
- h : step size
|
||||
'''
|
||||
# first steps is forward diff (t+1 - t) / dt
|
||||
init_vel = (data_seq[:, 1:2] - data_seq[:, :1]) / dt
|
||||
# middle steps are second order (t+1 - t-1) / 2dt
|
||||
middle_vel = (data_seq[:, 2:] - data_seq[:, 0:-2]) / (2 * dt)
|
||||
# last step is backward diff (t - t-1) / dt
|
||||
final_vel = (data_seq[:, -1:] - data_seq[:, -2:-1]) / dt
|
||||
|
||||
vel_seq = torch.cat([init_vel, middle_vel, final_vel], dim=1)
|
||||
return vel_seq
|
||||
|
||||
|
||||
def estimate_angular_velocity(rot_seq, dt):
|
||||
'''
|
||||
Given a batch of sequences of T rotation matrices, estimates angular velocity at T-2 steps.
|
||||
Input sequence should be of shape (B, T, ..., 3, 3)
|
||||
'''
|
||||
# see https://en.wikipedia.org/wiki/Angular_velocity#Calculation_from_the_orientation_matrix
|
||||
dRdt = estimate_linear_velocity(rot_seq, dt)
|
||||
R = rot_seq
|
||||
RT = R.transpose(-1, -2)
|
||||
# compute skew-symmetric angular velocity tensor
|
||||
w_mat = torch.matmul(dRdt, RT)
|
||||
# pull out angular velocity vector by averaging symmetric entries
|
||||
w_x = (-w_mat[..., 1, 2] + w_mat[..., 2, 1]) / 2.0
|
||||
w_y = (w_mat[..., 0, 2] - w_mat[..., 2, 0]) / 2.0
|
||||
w_z = (-w_mat[..., 0, 1] + w_mat[..., 1, 0]) / 2.0
|
||||
w = torch.stack([w_x, w_y, w_z], axis=-1)
|
||||
return w
|
||||
|
||||
|
||||
def normalize(tensor, mean=None, std=None):
|
||||
"""
|
||||
Args:
|
||||
tensor: (B, T, ...)
|
||||
|
||||
Returns:
|
||||
normalized tensor with 0 mean and 1 standard deviation, std, mean
|
||||
"""
|
||||
if mean is None or std is None:
|
||||
# std, mean = torch.std_mean(tensor, dim=0, unbiased=False, keepdim=True)
|
||||
std, mean = torch.std_mean(tensor, dim=(0, 1), unbiased=False, keepdim=True)
|
||||
std[std == 0.0] = 1.0
|
||||
|
||||
return (tensor - mean) / std, mean, std
|
||||
|
||||
return (tensor - mean) / std
|
||||
|
||||
|
||||
def denormalize(tensor, mean, std):
|
||||
"""
|
||||
Args:
|
||||
tensor: B x T x D
|
||||
mean:
|
||||
std:
|
||||
"""
|
||||
return tensor * std + mean
|
||||
|
||||
|
||||
def export_bvh_animation(rotations, positions, offsets, parents, output_dir, prefix, joint_names, fps):
|
||||
"""
|
||||
Args:
|
||||
rotations: quaternions of the shape (B, T, J, 4)
|
||||
positions: global translations of the shape (B, T, 3)
|
||||
"""
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
for i in range(rotations.shape[0]):
|
||||
rotation = align_joints(rotations[i])
|
||||
position = positions[i]
|
||||
position = position.unsqueeze(1)
|
||||
anim = Animation(Quaternions(c2c(rotation)), c2c(position), None, offsets=offsets, parents=parents)
|
||||
BVH.save(os.path.join(output_dir, f'{prefix}_{i}.bvh'), anim, names=joint_names, frametime=1 / fps)
|
||||
|
||||
|
||||
def export_ply_trajectory(points, color, ply_fname):
|
||||
v = []
|
||||
for p in points:
|
||||
v += [(p[0], p[1], p[2], color[0], color[1], color[2])]
|
||||
v = np.array(v, dtype=[('x', 'f4'), ('y', 'f4'), ('z', 'f4'), ('red', 'u1'), ('green', 'u1'), ('blue', 'u1')])
|
||||
el_v = plyfile.PlyElement.describe(v, 'vertex')
|
||||
|
||||
e = np.empty(len(points) - 1, dtype=[('vertex1', 'i4'), ('vertex2', 'i4')])
|
||||
edge_data = np.array([[i, i + 1] for i in range(len(points) - 1)], dtype='i4')
|
||||
e['vertex1'] = edge_data[:, 0]
|
||||
e['vertex2'] = edge_data[:, 1]
|
||||
el_e = plyfile.PlyElement.describe(e, 'edge')
|
||||
|
||||
plyfile.PlyData([el_v, el_e], text=True).write(ply_fname)
|
||||
|
||||
|
||||
def get_keypoints_rectangle(keypoints: np.array, threshold: float) -> Tuple[float, float, float]:
|
||||
"""
|
||||
Compute rectangle enclosing keypoints above the threshold.
|
||||
Args:
|
||||
keypoints (np.array): Keypoint array of shape (N, 3).
|
||||
threshold (float): Confidence visualization threshold.
|
||||
Returns:
|
||||
Tuple[float, float, float]: Rectangle width, height and area.
|
||||
"""
|
||||
# print(keypoints.shape)
|
||||
valid_ind = keypoints[:, -1] > threshold
|
||||
if valid_ind.sum() > 0:
|
||||
valid_keypoints = keypoints[valid_ind][:, :-1]
|
||||
max_x = valid_keypoints[:,0].max()
|
||||
max_y = valid_keypoints[:,1].max()
|
||||
min_x = valid_keypoints[:,0].min()
|
||||
min_y = valid_keypoints[:,1].min()
|
||||
width = max_x - min_x
|
||||
height = max_y - min_y
|
||||
area = width * height
|
||||
return width, height, area
|
||||
else:
|
||||
return 0,0,0
|
||||
|
||||
def render_keypoints(img: np.array,
|
||||
keypoints: np.array,
|
||||
pairs: List,
|
||||
colors: List,
|
||||
thickness_circle_ratio: float,
|
||||
thickness_line_ratio_wrt_circle: float,
|
||||
pose_scales: List,
|
||||
threshold: float = 0.1,
|
||||
alpha: float = 1.0) -> np.array:
|
||||
"""
|
||||
Render keypoints on input image.
|
||||
Args:
|
||||
img (np.array): Input image of shape (H, W, 3) with pixel values in the [0,255] range.
|
||||
keypoints (np.array): Keypoint array of shape (N, 3).
|
||||
pairs (List): List of keypoint pairs per limb.
|
||||
colors: (List): List of colors per keypoint.
|
||||
thickness_circle_ratio (float): Circle thickness ratio.
|
||||
thickness_line_ratio_wrt_circle (float): Line thickness ratio wrt the circle.
|
||||
pose_scales (List): List of pose scales.
|
||||
threshold (float): Only visualize keypoints with confidence above the threshold.
|
||||
Returns:
|
||||
(np.array): Image of shape (H, W, 3) with keypoints drawn on top of the original image.
|
||||
"""
|
||||
img_orig = img.copy()
|
||||
width, height = img.shape[1], img.shape[2]
|
||||
area = width * height
|
||||
|
||||
lineType = 8
|
||||
shift = 0
|
||||
numberColors = len(colors)
|
||||
thresholdRectangle = 0.1
|
||||
|
||||
person_width, person_height, person_area = get_keypoints_rectangle(keypoints, thresholdRectangle)
|
||||
if person_area > 0:
|
||||
ratioAreas = min(1, max(person_width / width, person_height / height))
|
||||
thicknessRatio = np.maximum(np.round(math.sqrt(area) * thickness_circle_ratio * ratioAreas), 2)
|
||||
thicknessCircle = np.maximum(1, thicknessRatio if ratioAreas > 0.05 else -np.ones_like(thicknessRatio))
|
||||
thicknessLine = np.maximum(1, np.round(thicknessRatio * thickness_line_ratio_wrt_circle))
|
||||
radius = thicknessRatio / 2
|
||||
|
||||
img = np.ascontiguousarray(img.copy())
|
||||
for i, pair in enumerate(pairs):
|
||||
index1, index2 = pair
|
||||
if keypoints[index1, -1] > threshold and keypoints[index2, -1] > threshold:
|
||||
thicknessLineScaled = int(round(min(thicknessLine[index1], thicknessLine[index2]) * pose_scales[0]))
|
||||
colorIndex = index2
|
||||
color = colors[colorIndex % numberColors]
|
||||
keypoint1 = keypoints[index1, :-1].astype(np.int32)
|
||||
keypoint2 = keypoints[index2, :-1].astype(np.int32)
|
||||
cv2.line(img, tuple(keypoint1.tolist()), tuple(keypoint2.tolist()), tuple(color.tolist()), thicknessLineScaled, lineType, shift)
|
||||
for part in range(len(keypoints)):
|
||||
faceIndex = part
|
||||
if keypoints[faceIndex, -1] > threshold:
|
||||
radiusScaled = int(round(radius[faceIndex] * pose_scales[0]))
|
||||
thicknessCircleScaled = int(round(thicknessCircle[faceIndex] * pose_scales[0]))
|
||||
colorIndex = part
|
||||
color = colors[colorIndex % numberColors]
|
||||
center = keypoints[faceIndex, :-1].astype(np.int32)
|
||||
cv2.circle(img, tuple(center.tolist()), radiusScaled, tuple(color.tolist()), thicknessCircleScaled, lineType, shift)
|
||||
return img
|
||||
|
||||
def render_hand_keypoints(img, right_hand_keypoints, threshold=0.1, use_confidence=False, map_fn=lambda x: np.ones_like(x), alpha=1.0):
|
||||
if use_confidence and map_fn is not None:
|
||||
#thicknessCircleRatioLeft = 1./50 * map_fn(left_hand_keypoints[:, -1])
|
||||
thicknessCircleRatioRight = 1./50 * map_fn(right_hand_keypoints[:, -1])
|
||||
else:
|
||||
#thicknessCircleRatioLeft = 1./50 * np.ones(left_hand_keypoints.shape[0])
|
||||
thicknessCircleRatioRight = 1./50 * np.ones(right_hand_keypoints.shape[0])
|
||||
thicknessLineRatioWRTCircle = 0.75
|
||||
pairs = [0,1, 1,2, 2,3, 3,4, 0,5, 5,6, 6,7, 7,8, 0,9, 9,10, 10,11, 11,12, 0,13, 13,14, 14,15, 15,16, 0,17, 17,18, 18,19, 19,20]
|
||||
pairs = np.array(pairs).reshape(-1,2)
|
||||
|
||||
colors = [100., 100., 100.,
|
||||
100., 0., 0.,
|
||||
150., 0., 0.,
|
||||
200., 0., 0.,
|
||||
255., 0., 0.,
|
||||
100., 100., 0.,
|
||||
150., 150., 0.,
|
||||
200., 200., 0.,
|
||||
255., 255., 0.,
|
||||
0., 100., 50.,
|
||||
0., 150., 75.,
|
||||
0., 200., 100.,
|
||||
0., 255., 125.,
|
||||
0., 50., 100.,
|
||||
0., 75., 150.,
|
||||
0., 100., 200.,
|
||||
0., 125., 255.,
|
||||
100., 0., 100.,
|
||||
150., 0., 150.,
|
||||
200., 0., 200.,
|
||||
255., 0., 255.]
|
||||
colors = np.array(colors).reshape(-1,3)
|
||||
#colors = np.zeros_like(colors)
|
||||
poseScales = [1]
|
||||
#img = render_keypoints(img, left_hand_keypoints, pairs, colors, thicknessCircleRatioLeft, thicknessLineRatioWRTCircle, poseScales, threshold, alpha=alpha)
|
||||
img = render_keypoints(img, right_hand_keypoints, pairs, colors, thicknessCircleRatioRight, thicknessLineRatioWRTCircle, poseScales, threshold, alpha=alpha)
|
||||
#img = render_keypoints(img, right_hand_keypoints, pairs, colors, thickness_circle_ratio, thickness_line_ratio_wrt_circle, pose_scales, 0.1)
|
||||
return img
|
||||
|
||||
def render_body_keypoints(img: np.array,
|
||||
body_keypoints: np.array) -> np.array:
|
||||
"""
|
||||
Render OpenPose body keypoints on input image.
|
||||
Args:
|
||||
img (np.array): Input image of shape (H, W, 3) with pixel values in the [0,255] range.
|
||||
body_keypoints (np.array): Keypoint array of shape (N, 3); 3 <====> (x, y, confidence).
|
||||
Returns:
|
||||
(np.array): Image of shape (H, W, 3) with keypoints drawn on top of the original image.
|
||||
"""
|
||||
|
||||
thickness_circle_ratio = 1./75. * np.ones(body_keypoints.shape[0])
|
||||
thickness_line_ratio_wrt_circle = 0.75
|
||||
pairs = []
|
||||
pairs = [1,8,1,2,1,5,2,3,3,4,5,6,6,7,8,9,9,10,10,11,8,12,12,13,13,14,1,0,0,15,15,17,0,16,16,18,14,19,19,20,14,21,11,22,22,23,11,24]
|
||||
pairs = np.array(pairs).reshape(-1,2)
|
||||
colors = [255., 0., 85.,
|
||||
255., 0., 0.,
|
||||
255., 85., 0.,
|
||||
255., 170., 0.,
|
||||
255., 255., 0.,
|
||||
170., 255., 0.,
|
||||
85., 255., 0.,
|
||||
0., 255., 0.,
|
||||
255., 0., 0.,
|
||||
0., 255., 85.,
|
||||
0., 255., 170.,
|
||||
0., 255., 255.,
|
||||
0., 170., 255.,
|
||||
0., 85., 255.,
|
||||
0., 0., 255.,
|
||||
255., 0., 170.,
|
||||
170., 0., 255.,
|
||||
255., 0., 255.,
|
||||
85., 0., 255.,
|
||||
0., 0., 255.,
|
||||
0., 0., 255.,
|
||||
0., 0., 255.,
|
||||
0., 255., 255.,
|
||||
0., 255., 255.,
|
||||
0., 255., 255.]
|
||||
colors = np.array(colors).reshape(-1,3)
|
||||
pose_scales = [1]
|
||||
return render_keypoints(img, body_keypoints, pairs, colors, thickness_circle_ratio, thickness_line_ratio_wrt_circle, pose_scales, 0.1)
|
||||
|
||||
def render_openpose(img: np.array,
|
||||
hand_keypoints: np.array) -> np.array:
|
||||
"""
|
||||
Render keypoints in the OpenPose format on input image.
|
||||
Args:
|
||||
img (np.array): Input image of shape (H, W, 3) with pixel values in the [0,255] range.
|
||||
body_keypoints (np.array): Keypoint array of shape (N, 3); 3 <====> (x, y, confidence).
|
||||
Returns:
|
||||
(np.array): Image of shape (H, W, 3) with keypoints drawn on top of the original image.
|
||||
"""
|
||||
#img = render_body_keypoints(img, body_keypoints)
|
||||
img = render_hand_keypoints(img, hand_keypoints)
|
||||
return img
|
||||
+124
@@ -0,0 +1,124 @@
|
||||
"""Run the fixed 128-frame HMP on overlapping windows of a full sequence."""
|
||||
import copy
|
||||
import json
|
||||
from pathlib import Path
|
||||
import numpy as np
|
||||
import torch
|
||||
from scipy.spatial.transform import Rotation
|
||||
|
||||
|
||||
def window_starts(length, size=128, overlap=16):
|
||||
if length < 1 or not 0 <= overlap < size:
|
||||
raise ValueError('Invalid sequence length or overlap')
|
||||
starts = [0]
|
||||
while starts[-1] + size < length:
|
||||
starts.append(starts[-1] + size - overlap)
|
||||
return starts
|
||||
|
||||
|
||||
def slice_window(data, start, length, size=128):
|
||||
result = {}
|
||||
for key, value in data.items():
|
||||
if isinstance(value, (torch.Tensor, np.ndarray)) and value.ndim >= 2 and value.shape[1] == length:
|
||||
clip = value[:, start:min(start + size, length)].clone() if isinstance(value, torch.Tensor) else value[:, start:min(start + size, length)].copy()
|
||||
valid = clip.shape[1]
|
||||
if valid < size:
|
||||
if isinstance(clip, torch.Tensor):
|
||||
pad = clip[:, -1:].repeat_interleave(size-valid, dim=1)
|
||||
clip = torch.cat([clip, pad], dim=1)
|
||||
else:
|
||||
clip = np.concatenate([clip, np.repeat(clip[:, -1:], size-valid, axis=1)], axis=1)
|
||||
if key == 'vis_mask':
|
||||
clip[:, valid:] = 0
|
||||
elif key == 'joints2d':
|
||||
clip[:, valid:, :, -1] = 0
|
||||
result[key] = clip
|
||||
else:
|
||||
result[key] = copy.deepcopy(value)
|
||||
return result
|
||||
|
||||
|
||||
def blend_windows(original, clips, size=128, overlap=16):
|
||||
result = {k: np.array(v, copy=True) for k, v in original.items()}
|
||||
B, T = result['trans'].shape[:2]
|
||||
denom = np.zeros(T, dtype=np.float64)
|
||||
trans = np.zeros((B, T, 3), dtype=np.float64)
|
||||
rotation_keys = ['root_orient', 'pose_body']
|
||||
if all('decode_root' in clip for _, clip in clips):
|
||||
result['decode_root'] = np.zeros((B, T, 3), dtype=np.float32)
|
||||
rotation_keys.append('decode_root')
|
||||
accum, refs = {}, {}
|
||||
for key in rotation_keys:
|
||||
shape = result[key].shape
|
||||
refs[key] = Rotation.from_rotvec(result[key].reshape(-1, 3)).as_quat().reshape(*shape[:-1], 4)
|
||||
accum[key] = np.zeros_like(refs[key])
|
||||
for start, clip in clips:
|
||||
n = min(size, T-start)
|
||||
w = np.ones(n, dtype=np.float64)
|
||||
if start > 0:
|
||||
m = min(overlap, n)
|
||||
w[:m] *= np.arange(1, m+1)/(m+1)
|
||||
if start + size < T:
|
||||
m = min(overlap, n)
|
||||
w[-m:] *= np.arange(m, 0, -1)/(m+1)
|
||||
denom[start:start+n] += w
|
||||
trans[:, start:start+n] += clip['trans'][:, :n] * w[None, :, None]
|
||||
for key in rotation_keys:
|
||||
a = clip[key][:, :n].reshape(result[key][:, start:start+n].shape)
|
||||
q = Rotation.from_rotvec(a.reshape(-1, 3)).as_quat().reshape(*a.shape[:-1], 4)
|
||||
reference = refs[key][:, start:start+n]
|
||||
q *= np.where(np.sum(q*reference, axis=-1, keepdims=True) < 0, -1., 1.)
|
||||
weights = w.reshape((1,n) + (1,)*(q.ndim-2))
|
||||
accum[key][:, start:start+n] += q * weights
|
||||
if np.any(denom <= 0):
|
||||
raise ValueError('Uncovered frames in HMP windows')
|
||||
result['trans'] = (trans/denom[None,:,None]).astype(np.float32)
|
||||
for key, q in accum.items():
|
||||
if np.any(np.linalg.norm(q, axis=-1) < 1e-8):
|
||||
raise ValueError('Degenerate quaternion blend')
|
||||
result[key] = Rotation.from_quat(q.reshape(-1,4)).as_rotvec().reshape(result[key].shape).astype(np.float32)
|
||||
# HMP currently fixes shape. Reject unhandled per-window shape changes.
|
||||
for _, clip in clips:
|
||||
if not np.allclose(clip['betas'], result['betas'], atol=1e-5):
|
||||
raise ValueError('Windowed HMP requires fixed shape (opt_betas=false)')
|
||||
result['latent_pose'] = result['pose_body'].copy()
|
||||
return result
|
||||
|
||||
|
||||
def fit_windowed(obs_data, res_dicts, opt, data_args, out_dir, device, fit_fn):
|
||||
from body_model import MANO
|
||||
original = {k: v.detach().cpu().numpy() if isinstance(v, torch.Tensor) else np.asarray(v) for k,v in res_dicts[0].items()}
|
||||
T = original['trans'].shape[1]
|
||||
starts = window_starts(T)
|
||||
out = Path(out_dir)
|
||||
out.mkdir(parents=True, exist_ok=True)
|
||||
mano_cfg = {k.lower():v for k,v in dict(opt.MANO).items()}
|
||||
hand_model = MANO(batch_size=128, pose2rot=True, **mano_cfg).to(device)
|
||||
clips = []
|
||||
manifest = {'frames':T, 'window_size':128, 'overlap':16, 'windows':[]}
|
||||
for index, start in enumerate(starts):
|
||||
clip_dir = out / 'windows' / f'{start:06d}'
|
||||
path = clip_dir / f'{Path(opt.HMP.vid_path).stem}_000000_world_results.npz'
|
||||
print(f'HMP window {index+1}/{len(starts)}: frames {start}:{min(start+128,T)}', flush=True)
|
||||
if path.is_file():
|
||||
with np.load(path) as cached:
|
||||
compatible = np.allclose(cached['betas'], original['betas'], atol=1e-5)
|
||||
if not compatible:
|
||||
path.rename(path.with_suffix('.variable_shape_backup.npz'))
|
||||
if not path.is_file():
|
||||
obs = slice_window(obs_data, start, T)
|
||||
res = slice_window(res_dicts[0], start, T)
|
||||
fit_fn(obs, (res,), hand_model, opt, data_args, str(clip_dir), device)
|
||||
with np.load(path) as saved:
|
||||
clip = {k:saved[k].copy() for k in saved.files}
|
||||
if clip['trans'].shape[1] != 128 or not all(np.isfinite(v).all() for v in clip.values() if v.dtype.kind in 'fc'):
|
||||
raise ValueError(f'Invalid HMP window output: {path}')
|
||||
clips.append((start,clip))
|
||||
manifest['windows'].append({'start':start,'valid_frames':min(128,T-start),'path':str(path)})
|
||||
(out/'window_manifest.json').write_text(json.dumps(manifest,indent=2)+'\n')
|
||||
result = blend_windows(original,clips)
|
||||
target = out/f'{Path(opt.HMP.vid_path).stem}_000000_world_results.npz'
|
||||
temporary = target.with_suffix('.partial.npz')
|
||||
np.savez(temporary,**result)
|
||||
temporary.replace(target)
|
||||
print(f'Saved full-length windowed HMP: {target}', flush=True)
|
||||
Reference in New Issue
Block a user