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)
|
||||
@@ -0,0 +1,4 @@
|
||||
from .body_model import *
|
||||
from .specs import *
|
||||
from .utils import *
|
||||
from .mano_wrapper import *
|
||||
+2415
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,44 @@
|
||||
import torch
|
||||
import numpy as np
|
||||
import pickle
|
||||
from typing import Optional
|
||||
from body_model import MANOLayer
|
||||
from smplx.lbs import vertices2joints
|
||||
from smplx.utils import MANOOutput, to_tensor
|
||||
from smplx.vertex_ids import vertex_ids
|
||||
|
||||
|
||||
class MANO(MANOLayer):
|
||||
def __init__(self, *args, joint_regressor_extra: Optional[str] = None, **kwargs):
|
||||
"""
|
||||
Extension of the official MANO implementation to support more joints.
|
||||
Args:
|
||||
Same as MANOLayer.
|
||||
joint_regressor_extra (str): Path to extra joint regressor.
|
||||
"""
|
||||
super(MANO, self).__init__(*args, **kwargs)
|
||||
mano_to_openpose = [0, 13, 14, 15, 16, 1, 2, 3, 17, 4, 5, 6, 18, 10, 11, 12, 19, 7, 8, 9, 20]
|
||||
|
||||
#2, 3, 5, 4, 1
|
||||
if joint_regressor_extra is not None:
|
||||
self.register_buffer('joint_regressor_extra', torch.tensor(pickle.load(open(joint_regressor_extra, 'rb'), encoding='latin1'), dtype=torch.float32))
|
||||
self.register_buffer('extra_joints_idxs', to_tensor(list(vertex_ids['mano'].values()), dtype=torch.long))
|
||||
self.register_buffer('joint_map', torch.tensor(mano_to_openpose, dtype=torch.long))
|
||||
self.pose2rot = kwargs.get('pose2rot')
|
||||
# print([attr for attr in dir(self) if not callable(getattr(self, attr)) and not attr.startswith("__")])
|
||||
# print(self.hand_mean, self.flat_hand_mean, self.num_pca_comps, self.use_pca)
|
||||
# exit()
|
||||
|
||||
def forward(self, *args, **kwargs) -> MANOOutput:
|
||||
"""
|
||||
Run forward pass. Same as MANO and also append an extra set of joints if joint_regressor_extra is specified.
|
||||
"""
|
||||
mano_output = super(MANO, self).forward(*args, **kwargs)
|
||||
extra_joints = torch.index_select(mano_output.vertices, 1, self.extra_joints_idxs)
|
||||
joints = torch.cat([mano_output.joints, extra_joints], dim=1)
|
||||
joints = joints[:, self.joint_map, :]
|
||||
if hasattr(self, 'joint_regressor_extra'):
|
||||
extra_joints = vertices2joints(self.joint_regressor_extra, mano_output.vertices)
|
||||
joints = torch.cat([joints, extra_joints], dim=1)
|
||||
mano_output.joints = joints
|
||||
return mano_output
|
||||
+575
@@ -0,0 +1,575 @@
|
||||
import numpy as np
|
||||
|
||||
# SMPL_JOINTS = {
|
||||
# "hips": 0,
|
||||
# "leftUpLeg": 1,
|
||||
# "rightUpLeg": 2,
|
||||
# "spine": 3,
|
||||
# "leftLeg": 4,
|
||||
# "rightLeg": 5,
|
||||
# "spine1": 6,
|
||||
# "leftFoot": 7,
|
||||
# "rightFoot": 8,
|
||||
# "spine2": 9,
|
||||
# "leftToeBase": 10,
|
||||
# "rightToeBase": 11,
|
||||
# "neck": 12,
|
||||
# "leftShoulder": 13,
|
||||
# "rightShoulder": 14,
|
||||
# "head": 15,
|
||||
# "leftArm": 16,
|
||||
# "rightArm": 17,
|
||||
# "leftForeArm": 18,
|
||||
# "rightForeArm": 19,
|
||||
# "leftHand": 20,
|
||||
# "rightHand": 21,
|
||||
# }
|
||||
MANO_JOINTS = {
|
||||
'wrist': 0,
|
||||
'index1': 1,
|
||||
'index2': 2,
|
||||
'index3': 3,
|
||||
'middle1': 4,
|
||||
'middle2': 5,
|
||||
'middle3': 6,
|
||||
'pinky1': 7,
|
||||
'pinky2': 8,
|
||||
'pinky3': 9,
|
||||
'ring1': 10,
|
||||
'ring2': 11,
|
||||
'ring3': 12,
|
||||
'thumb1': 13,
|
||||
'thumb2': 14,
|
||||
'thumb3': 15,
|
||||
}
|
||||
OP_IGNORE_JOINTS = None
|
||||
OP_NUM_JOINTS = 21
|
||||
|
||||
# SMPL_PARENTS = [
|
||||
# -1,
|
||||
# 0,
|
||||
# 0,
|
||||
# 0,
|
||||
# 1,
|
||||
# 2,
|
||||
# 3,
|
||||
# 4,
|
||||
# 5,
|
||||
# 6,
|
||||
# 7,
|
||||
# 8,
|
||||
# 9,
|
||||
# 12,
|
||||
# 12,
|
||||
# 12,
|
||||
# 13,
|
||||
# 14,
|
||||
# 16,
|
||||
# 17,
|
||||
# 18,
|
||||
# 19,
|
||||
# ]
|
||||
|
||||
SMPLH_PATH = "./body_models/smplh"
|
||||
SMPLX_PATH = "./body_models/smplx"
|
||||
SMPL_PATH = "./body_models/smpl"
|
||||
VPOSER_PATH = "./body_models/vposer_v1_0"
|
||||
|
||||
# chosen virtual mocap markers that are "keypoints" to work with
|
||||
# KEYPT_VERTS = [
|
||||
# 4404,
|
||||
# 920,
|
||||
# 3076,
|
||||
# 3169,
|
||||
# 823,
|
||||
# 4310,
|
||||
# 1010,
|
||||
# 1085,
|
||||
# 4495,
|
||||
# 4569,
|
||||
# 6615,
|
||||
# 3217,
|
||||
# 3313,
|
||||
# 6713,
|
||||
# 6785,
|
||||
# 3383,
|
||||
# 6607,
|
||||
# 3207,
|
||||
# 1241,
|
||||
# 1508,
|
||||
# 4797,
|
||||
# 4122,
|
||||
# 1618,
|
||||
# 1569,
|
||||
# 5135,
|
||||
# 5040,
|
||||
# 5691,
|
||||
# 5636,
|
||||
# 5404,
|
||||
# 2230,
|
||||
# 2173,
|
||||
# 2108,
|
||||
# 134,
|
||||
# 3645,
|
||||
# 6543,
|
||||
# 3123,
|
||||
# 3024,
|
||||
# 4194,
|
||||
# 1306,
|
||||
# 182,
|
||||
# 3694,
|
||||
# 4294,
|
||||
# 744,
|
||||
# ]
|
||||
|
||||
|
||||
# """
|
||||
# Openpose
|
||||
# """
|
||||
# OP_NUM_JOINTS = 25
|
||||
# # OP_IGNORE_JOINTS = [1, 9, 12] # neck and left/right hip
|
||||
# OP_IGNORE_JOINTS = [1] # neck
|
||||
# OP_EDGE_LIST = [
|
||||
# [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],
|
||||
# ]
|
||||
# # indices to map an openpose detection to its flipped version
|
||||
# OP_FLIP_MAP = [
|
||||
# 0,
|
||||
# 1,
|
||||
# 5,
|
||||
# 6,
|
||||
# 7,
|
||||
# 2,
|
||||
# 3,
|
||||
# 4,
|
||||
# 8,
|
||||
# 12,
|
||||
# 13,
|
||||
# 14,
|
||||
# 9,
|
||||
# 10,
|
||||
# 11,
|
||||
# 16,
|
||||
# 15,
|
||||
# 18,
|
||||
# 17,
|
||||
# 22,
|
||||
# 23,
|
||||
# 24,
|
||||
# 19,
|
||||
# 20,
|
||||
# 21,
|
||||
# ]
|
||||
|
||||
|
||||
#
|
||||
# From https://github.com/vchoutas/smplify-x/blob/master/smplifyx/utils.py
|
||||
# Please see license for usage restrictions.
|
||||
#
|
||||
# def smpl_to_openpose(
|
||||
# model_type="smplx",
|
||||
# use_hands=True,
|
||||
# use_face=True,
|
||||
# use_face_contour=False,
|
||||
# openpose_format="coco25",
|
||||
# ):
|
||||
# """Returns the indices of the permutation that maps SMPL to OpenPose
|
||||
|
||||
# Parameters
|
||||
# ----------
|
||||
# model_type: str, optional
|
||||
# The type of SMPL-like model that is used. The default mapping
|
||||
# returned is for the SMPLX model
|
||||
# use_hands: bool, optional
|
||||
# Flag for adding to the returned permutation the mapping for the
|
||||
# hand keypoints. Defaults to True
|
||||
# use_face: bool, optional
|
||||
# Flag for adding to the returned permutation the mapping for the
|
||||
# face keypoints. Defaults to True
|
||||
# use_face_contour: bool, optional
|
||||
# Flag for appending the facial contour keypoints. Defaults to False
|
||||
# openpose_format: bool, optional
|
||||
# The output format of OpenPose. For now only COCO-25 and COCO-19 is
|
||||
# supported. Defaults to 'coco25'
|
||||
|
||||
# """
|
||||
# if openpose_format.lower() == "coco25":
|
||||
# if model_type == "smpl":
|
||||
# return np.array(
|
||||
# [
|
||||
# 24,
|
||||
# 12,
|
||||
# 17,
|
||||
# 19,
|
||||
# 21,
|
||||
# 16,
|
||||
# 18,
|
||||
# 20,
|
||||
# 0,
|
||||
# 2,
|
||||
# 5,
|
||||
# 8,
|
||||
# 1,
|
||||
# 4,
|
||||
# 7,
|
||||
# 25,
|
||||
# 26,
|
||||
# 27,
|
||||
# 28,
|
||||
# 29,
|
||||
# 30,
|
||||
# 31,
|
||||
# 32,
|
||||
# 33,
|
||||
# 34,
|
||||
# ],
|
||||
# dtype=np.int32,
|
||||
# )
|
||||
# elif model_type == "smplh":
|
||||
# body_mapping = np.array(
|
||||
# [
|
||||
# 52,
|
||||
# 12,
|
||||
# 17,
|
||||
# 19,
|
||||
# 21,
|
||||
# 16,
|
||||
# 18,
|
||||
# 20,
|
||||
# 0,
|
||||
# 2,
|
||||
# 5,
|
||||
# 8,
|
||||
# 1,
|
||||
# 4,
|
||||
# 7,
|
||||
# 53,
|
||||
# 54,
|
||||
# 55,
|
||||
# 56,
|
||||
# 57,
|
||||
# 58,
|
||||
# 59,
|
||||
# 60,
|
||||
# 61,
|
||||
# 62,
|
||||
# ],
|
||||
# dtype=np.int32,
|
||||
# )
|
||||
# mapping = [body_mapping]
|
||||
# if use_hands:
|
||||
# lhand_mapping = np.array(
|
||||
# [
|
||||
# 20,
|
||||
# 34,
|
||||
# 35,
|
||||
# 36,
|
||||
# 63,
|
||||
# 22,
|
||||
# 23,
|
||||
# 24,
|
||||
# 64,
|
||||
# 25,
|
||||
# 26,
|
||||
# 27,
|
||||
# 65,
|
||||
# 31,
|
||||
# 32,
|
||||
# 33,
|
||||
# 66,
|
||||
# 28,
|
||||
# 29,
|
||||
# 30,
|
||||
# 67,
|
||||
# ],
|
||||
# dtype=np.int32,
|
||||
# )
|
||||
# rhand_mapping = np.array(
|
||||
# [
|
||||
# 21,
|
||||
# 49,
|
||||
# 50,
|
||||
# 51,
|
||||
# 68,
|
||||
# 37,
|
||||
# 38,
|
||||
# 39,
|
||||
# 69,
|
||||
# 40,
|
||||
# 41,
|
||||
# 42,
|
||||
# 70,
|
||||
# 46,
|
||||
# 47,
|
||||
# 48,
|
||||
# 71,
|
||||
# 43,
|
||||
# 44,
|
||||
# 45,
|
||||
# 72,
|
||||
# ],
|
||||
# dtype=np.int32,
|
||||
# )
|
||||
# mapping += [lhand_mapping, rhand_mapping]
|
||||
# return np.concatenate(mapping)
|
||||
# # SMPLX
|
||||
# elif model_type == "smplx":
|
||||
# body_mapping = np.array(
|
||||
# [
|
||||
# 55,
|
||||
# 12,
|
||||
# 17,
|
||||
# 19,
|
||||
# 21,
|
||||
# 16,
|
||||
# 18,
|
||||
# 20,
|
||||
# 0,
|
||||
# 2,
|
||||
# 5,
|
||||
# 8,
|
||||
# 1,
|
||||
# 4,
|
||||
# 7,
|
||||
# 56,
|
||||
# 57,
|
||||
# 58,
|
||||
# 59,
|
||||
# 60,
|
||||
# 61,
|
||||
# 62,
|
||||
# 63,
|
||||
# 64,
|
||||
# 65,
|
||||
# ],
|
||||
# dtype=np.int32,
|
||||
# )
|
||||
# mapping = [body_mapping]
|
||||
# if use_hands:
|
||||
# lhand_mapping = np.array(
|
||||
# [
|
||||
# 20,
|
||||
# 37,
|
||||
# 38,
|
||||
# 39,
|
||||
# 66,
|
||||
# 25,
|
||||
# 26,
|
||||
# 27,
|
||||
# 67,
|
||||
# 28,
|
||||
# 29,
|
||||
# 30,
|
||||
# 68,
|
||||
# 34,
|
||||
# 35,
|
||||
# 36,
|
||||
# 69,
|
||||
# 31,
|
||||
# 32,
|
||||
# 33,
|
||||
# 70,
|
||||
# ],
|
||||
# dtype=np.int32,
|
||||
# )
|
||||
# rhand_mapping = np.array(
|
||||
# [
|
||||
# 21,
|
||||
# 52,
|
||||
# 53,
|
||||
# 54,
|
||||
# 71,
|
||||
# 40,
|
||||
# 41,
|
||||
# 42,
|
||||
# 72,
|
||||
# 43,
|
||||
# 44,
|
||||
# 45,
|
||||
# 73,
|
||||
# 49,
|
||||
# 50,
|
||||
# 51,
|
||||
# 74,
|
||||
# 46,
|
||||
# 47,
|
||||
# 48,
|
||||
# 75,
|
||||
# ],
|
||||
# dtype=np.int32,
|
||||
# )
|
||||
|
||||
# mapping += [lhand_mapping, rhand_mapping]
|
||||
# if use_face:
|
||||
# # end_idx = 127 + 17 * use_face_contour
|
||||
# face_mapping = np.arange(
|
||||
# 76, 127 + 17 * use_face_contour, dtype=np.int32
|
||||
# )
|
||||
# mapping += [face_mapping]
|
||||
|
||||
# return np.concatenate(mapping)
|
||||
# else:
|
||||
# raise ValueError("Unknown model type: {}".format(model_type))
|
||||
# elif openpose_format == "coco19":
|
||||
# if model_type == "smpl":
|
||||
# return np.array(
|
||||
# [24, 12, 17, 19, 21, 16, 18, 20, 0, 2, 5, 8, 1, 4, 7, 25, 26, 27, 28],
|
||||
# dtype=np.int32,
|
||||
# )
|
||||
# elif model_type == "smplh":
|
||||
# body_mapping = np.array(
|
||||
# [52, 12, 17, 19, 21, 16, 18, 20, 0, 2, 5, 8, 1, 4, 7, 53, 54, 55, 56],
|
||||
# dtype=np.int32,
|
||||
# )
|
||||
# mapping = [body_mapping]
|
||||
# if use_hands:
|
||||
# lhand_mapping = np.array(
|
||||
# [
|
||||
# 20,
|
||||
# 34,
|
||||
# 35,
|
||||
# 36,
|
||||
# 57,
|
||||
# 22,
|
||||
# 23,
|
||||
# 24,
|
||||
# 58,
|
||||
# 25,
|
||||
# 26,
|
||||
# 27,
|
||||
# 59,
|
||||
# 31,
|
||||
# 32,
|
||||
# 33,
|
||||
# 60,
|
||||
# 28,
|
||||
# 29,
|
||||
# 30,
|
||||
# 61,
|
||||
# ],
|
||||
# dtype=np.int32,
|
||||
# )
|
||||
# rhand_mapping = np.array(
|
||||
# [
|
||||
# 21,
|
||||
# 49,
|
||||
# 50,
|
||||
# 51,
|
||||
# 62,
|
||||
# 37,
|
||||
# 38,
|
||||
# 39,
|
||||
# 63,
|
||||
# 40,
|
||||
# 41,
|
||||
# 42,
|
||||
# 64,
|
||||
# 46,
|
||||
# 47,
|
||||
# 48,
|
||||
# 65,
|
||||
# 43,
|
||||
# 44,
|
||||
# 45,
|
||||
# 66,
|
||||
# ],
|
||||
# dtype=np.int32,
|
||||
# )
|
||||
# mapping += [lhand_mapping, rhand_mapping]
|
||||
# return np.concatenate(mapping)
|
||||
# # SMPLX
|
||||
# elif model_type == "smplx":
|
||||
# body_mapping = np.array(
|
||||
# [55, 12, 17, 19, 21, 16, 18, 20, 0, 2, 5, 8, 1, 4, 7, 56, 57, 58, 59],
|
||||
# dtype=np.int32,
|
||||
# )
|
||||
# mapping = [body_mapping]
|
||||
# if use_hands:
|
||||
# lhand_mapping = np.array(
|
||||
# [
|
||||
# 20,
|
||||
# 37,
|
||||
# 38,
|
||||
# 39,
|
||||
# 60,
|
||||
# 25,
|
||||
# 26,
|
||||
# 27,
|
||||
# 61,
|
||||
# 28,
|
||||
# 29,
|
||||
# 30,
|
||||
# 62,
|
||||
# 34,
|
||||
# 35,
|
||||
# 36,
|
||||
# 63,
|
||||
# 31,
|
||||
# 32,
|
||||
# 33,
|
||||
# 64,
|
||||
# ],
|
||||
# dtype=np.int32,
|
||||
# )
|
||||
# rhand_mapping = np.array(
|
||||
# [
|
||||
# 21,
|
||||
# 52,
|
||||
# 53,
|
||||
# 54,
|
||||
# 65,
|
||||
# 40,
|
||||
# 41,
|
||||
# 42,
|
||||
# 66,
|
||||
# 43,
|
||||
# 44,
|
||||
# 45,
|
||||
# 67,
|
||||
# 49,
|
||||
# 50,
|
||||
# 51,
|
||||
# 68,
|
||||
# 46,
|
||||
# 47,
|
||||
# 48,
|
||||
# 69,
|
||||
# ],
|
||||
# dtype=np.int32,
|
||||
# )
|
||||
|
||||
# mapping += [lhand_mapping, rhand_mapping]
|
||||
# if use_face:
|
||||
# face_mapping = np.arange(
|
||||
# 70, 70 + 51 + 17 * use_face_contour, dtype=np.int32
|
||||
# )
|
||||
# mapping += [face_mapping]
|
||||
|
||||
# return np.concatenate(mapping)
|
||||
# else:
|
||||
# raise ValueError("Unknown model type: {}".format(model_type))
|
||||
# else:
|
||||
# raise ValueError("Unknown joint format: {}".format(openpose_format))
|
||||
@@ -0,0 +1,76 @@
|
||||
import torch
|
||||
from .specs import MANO_JOINTS
|
||||
|
||||
|
||||
def run_mano(body_model, trans, root_orient, body_pose, is_right, betas=None):
|
||||
"""
|
||||
Forward pass of the MANO model and populates pred_data accordingly with
|
||||
joints3d, verts3d, points3d.
|
||||
|
||||
trans : B x T x 3
|
||||
root_orient : B x T x 3
|
||||
body_pose : B x T x J*3
|
||||
betas : (optional) B x D
|
||||
"""
|
||||
B, T, _ = trans.shape
|
||||
bm_batch_size = body_model.batch_size
|
||||
assert bm_batch_size % B == 0
|
||||
# assert (is_right[0]==is_right[0][0]).all(), f'{is_right}'
|
||||
# assert (is_right[1]==is_right[1][0]).all(), f'{is_right}'
|
||||
|
||||
seq_len = bm_batch_size // B
|
||||
bm_num_betas = body_model.num_betas
|
||||
J_BODY = len(MANO_JOINTS) - 1 # all joints except root
|
||||
# print('utils.py, seq_len, T: ', seq_len, T, bm_batch_size, B) # 1, 120, 1, 1
|
||||
# print('trans, root_orient, body_pose:', trans.shape, root_orient.shape, body_pose.shape)
|
||||
if T == 1:
|
||||
raise ValueError
|
||||
# must expand to use with body model
|
||||
trans = trans.expand(B, seq_len, 3)
|
||||
root_orient = root_orient.expand(B, seq_len, 3)
|
||||
body_pose = body_pose.expand(B, seq_len, J_BODY * 3)
|
||||
elif T != seq_len:
|
||||
trans, root_orient, body_pose = zero_pad_tensors(
|
||||
[trans, root_orient, body_pose], seq_len - T
|
||||
)
|
||||
if betas is None:
|
||||
betas = torch.zeros(B, bm_num_betas, device=trans.device)
|
||||
betas = betas.reshape((B, 1, bm_num_betas)).expand((B, seq_len, bm_num_betas))
|
||||
# print('body_pose: ', body_pose.reshape((B * seq_len, -1)).shape)
|
||||
|
||||
mano_output = body_model(
|
||||
hand_pose=body_pose.reshape((B * seq_len, -1)),
|
||||
betas=betas.reshape((B * seq_len, -1)),
|
||||
global_orient=root_orient.reshape((B * seq_len, -1)),
|
||||
transl=trans.reshape((B * seq_len, -1)),
|
||||
)
|
||||
joints = mano_output.joints
|
||||
verts = mano_output.vertices
|
||||
|
||||
joints = joints.reshape(B, seq_len, -1, 3)[:, :T]
|
||||
verts = verts.reshape(B, seq_len, -1, 3)[:, :T]
|
||||
|
||||
is_right = is_right.unsqueeze(-1)
|
||||
joints[:, :, :, 0] = (2*is_right-1)*joints[:, :, :, 0]
|
||||
verts[:, :, :, 0] = (2*is_right-1)*verts[:, :, :, 0]
|
||||
|
||||
return {
|
||||
"joints": joints,
|
||||
"vertices": verts,
|
||||
"l_faces": body_model.faces_tensor[:,[0,2,1]],
|
||||
"r_faces": body_model.faces_tensor,
|
||||
"is_right": is_right.squeeze(-1),
|
||||
'body_pose': body_pose.clone()
|
||||
}
|
||||
|
||||
|
||||
def zero_pad_tensors(pad_list, pad_size):
|
||||
"""
|
||||
Assumes tensors in pad_list are B x T x D and pad temporal dimension
|
||||
"""
|
||||
B = pad_list[0].size(0)
|
||||
new_pad_list = []
|
||||
for pad_idx, pad_tensor in enumerate(pad_list):
|
||||
padding = torch.zeros((B, pad_size, pad_tensor.size(2))).to(pad_tensor)
|
||||
new_pad_list.append(torch.cat([pad_tensor, padding], dim=1))
|
||||
return new_pad_list
|
||||
+64
@@ -0,0 +1,64 @@
|
||||
defaults:
|
||||
- data: HOT3D
|
||||
- optim
|
||||
- _self_
|
||||
|
||||
model:
|
||||
use_init: True
|
||||
opt_cams: False #True
|
||||
opt_scale: True
|
||||
async_tracks: True
|
||||
|
||||
overwrite: False
|
||||
run_opt: True
|
||||
run_vis: True
|
||||
run_prior: False
|
||||
datatype: ${data.video_dir}
|
||||
vis:
|
||||
phases:
|
||||
# - motion_chunks
|
||||
# - input
|
||||
# - root_fit
|
||||
- smooth_fit
|
||||
- prior
|
||||
# - init
|
||||
render_views:
|
||||
- above
|
||||
- side
|
||||
- front
|
||||
- src_cam
|
||||
make_grid: True
|
||||
overwrite: True
|
||||
render_keypoints: False # Overlay GT 2D keypoints on source camera view
|
||||
|
||||
MANO:
|
||||
MODEL_PATH: ${paths.DATA_DIR}/mano
|
||||
GENDER: neutral
|
||||
NUM_HAND_JOINTS: 15
|
||||
MEAN_PARAMS: ${paths.DATA_DIR}/mano_mean_params.npz
|
||||
CREATE_BODY_POSE: FALSE
|
||||
|
||||
paths:
|
||||
base_dir: None
|
||||
DATA_DIR: _DATA/data/
|
||||
vposer: VPoser/pretrained/Vposer_right_mirrored # _DATA/body_models/vposer_v1_0
|
||||
|
||||
HMP:
|
||||
config: hmp_config.yaml
|
||||
use_hposer: False
|
||||
vid_path: ${data.root}/video/${data.seq}.mp4
|
||||
exp_name: None
|
||||
|
||||
fps: 30
|
||||
log_root: ../outputs/logs
|
||||
log_dir: ${log_root}/${data.type}-${data.split}
|
||||
exp_name: ${now:%Y-%m-%d}
|
||||
temporal_smooth: True
|
||||
gpu: 0
|
||||
is_static: False
|
||||
|
||||
hydra:
|
||||
job:
|
||||
chdir: True
|
||||
run:
|
||||
dir: ${log_dir}/${exp_name}/${data.name}
|
||||
@@ -0,0 +1,24 @@
|
||||
type: video
|
||||
split: custom
|
||||
root: /data/home/zy3023/code/hand/Dyn-HaMR/test # put your videos in root/videos/vid.mp4
|
||||
video_dir: videos
|
||||
seq: demo1
|
||||
ext: mp4
|
||||
src_path: ${data.root}/${data.video_dir}/${data.seq}.${data.ext}
|
||||
frame_opts:
|
||||
ext: jpg
|
||||
fps: 25
|
||||
start_sec: 0
|
||||
end_sec: -1
|
||||
use_cams: True
|
||||
track_ids: "all"
|
||||
shot_idx: 0
|
||||
start_idx: 0
|
||||
end_idx: 500
|
||||
split_cameras: True
|
||||
name: ${data.seq}-${data.track_ids}-shot-${data.shot_idx}-${data.start_idx}-${data.end_idx}
|
||||
sources:
|
||||
images: ${data.root}/images/${data.seq}
|
||||
cameras: ${data.root}/dynhamr/cameras/${data.seq}/shot-${data.shot_idx}
|
||||
tracks: ${data.root}/dynhamr/track_preds/${data.seq}
|
||||
shots: ${data.root}/dynhamr/shot_idcs/${data.seq}.json
|
||||
@@ -0,0 +1,29 @@
|
||||
type: video
|
||||
split: custom
|
||||
root: /data/home/zy3023/code/hand/Dyn-HaMR/test # put your videos in root/videos/vid.mp4
|
||||
video_dir: videos
|
||||
seq: prod1
|
||||
ext: mp4
|
||||
src_path: ${data.root}/${data.video_dir}/${data.seq}.${data.ext}
|
||||
frame_opts:
|
||||
ext: jpg
|
||||
fps: 25
|
||||
start_sec: 0
|
||||
end_sec: -1
|
||||
use_cams: True
|
||||
track_ids: "all"
|
||||
shot_idx: 0
|
||||
start_idx: 0
|
||||
end_idx: -1
|
||||
split_cameras: True
|
||||
name: ${data.seq}-${data.track_ids}-shot-${data.shot_idx}-${data.start_idx}-${data.end_idx}
|
||||
sources:
|
||||
images: ${data.root}/images/${data.seq}
|
||||
cameras: ${data.root}/dynhamr/cameras/${data.seq}/shot-${data.shot_idx}
|
||||
tracks: ${data.root}/dynhamr/track_preds/${data.seq}
|
||||
shots: ${data.root}/dynhamr/shot_idcs/${data.seq}.json
|
||||
|
||||
# VIPE-specific options
|
||||
use_vipe: True
|
||||
vipe_dir: /data/home/zy3023/code/hand/Dyn-HaMR/third-party/vipe/vipe_results
|
||||
|
||||
+13
@@ -0,0 +1,13 @@
|
||||
defaults:
|
||||
- data: posetrack
|
||||
- _self_
|
||||
|
||||
gap: 1
|
||||
log_root: outputs
|
||||
save_per_frame: False
|
||||
print_err: False
|
||||
stride: 48
|
||||
|
||||
hydra:
|
||||
run:
|
||||
dir: ${log_root}/init/${data.type}-${data.seq}-${data.depth_dir}-${data.fov}fov-gap${gap}
|
||||
+52
@@ -0,0 +1,52 @@
|
||||
optim:
|
||||
options:
|
||||
robust_loss_type: "bisquare"
|
||||
robust_tuning_const: 4.6851
|
||||
joints2d_sigma: 100.0
|
||||
lr: 1.0
|
||||
lbfgs_max_iter: 20
|
||||
save_every: 20
|
||||
vis_every: -1
|
||||
max_chunk_steps: 20
|
||||
save_meshes: False
|
||||
|
||||
root:
|
||||
num_iters: 50
|
||||
|
||||
smpl:
|
||||
num_iters: 0
|
||||
|
||||
smooth:
|
||||
opt_scale: False
|
||||
num_iters: 300
|
||||
|
||||
motion_chunks:
|
||||
chunk_size: 10
|
||||
init_steps: 20
|
||||
chunk_steps: 20
|
||||
opt_cams: True
|
||||
|
||||
loss_weights:
|
||||
penetration: [0,0,0] # [0, 1.0, 10.0]
|
||||
# joints2d: [1, 1, 1] # Increased smooth_fit weight from 0.001 to 0.1 to better match 2D keypoints
|
||||
joints2d: [10000, 10000, 10000] # Increased smooth_fit weight from 0.001 to 0.1 to better match 2D keypoints
|
||||
bg2d: [0.0, 0.000, 0.000]
|
||||
cam_R_smooth : [0.0, 0.0, 0.0]
|
||||
cam_t_smooth : [0.0, 0.0, 0.0]
|
||||
bio : [0.0, 0.0, 0.0] # [10, 10, 10]
|
||||
depth_constraint: [100.0, 100.0, 0.0] # Prevent hand from being behind camera - MUST be stronger than 2D loss
|
||||
# bg2d: [0.0, 0.0001, 0.0001]
|
||||
# cam_R_smooth : [0.0, 1000.0, 1000.0]
|
||||
# cam_t_smooth : [0.0, 1000.0, 1000.0]
|
||||
joints3d: [0.0, 0.0, 0.0]
|
||||
joints3d_smooth: [1000, 10000.0, 0.0] # raw: [1.0, 10.0, 0.0]
|
||||
# joints3d_smooth: [0, 0, 0.0] # raw: [1.0, 10.0, 0.0]
|
||||
joints3d_rollout: [0.0, 0.0, 0.0]
|
||||
verts3d: [0.0, 0.0, 0.0]
|
||||
points3d: [0.0, 0.0, 0.0]
|
||||
pose_prior: [1, 1, 1]
|
||||
shape_prior: [0.05, 0.05, 0.05]
|
||||
motion_prior: [0.0, 0.0, 0.075]
|
||||
init_motion_prior: [0.0, 0.0, 0.075]
|
||||
joint_consistency: [0.0, 0.0, 100.0]
|
||||
bone_length: [0.0, 2000.0, 2000.0]
|
||||
@@ -0,0 +1,2 @@
|
||||
from .dataset import *
|
||||
from . import tools
|
||||
+482
@@ -0,0 +1,482 @@
|
||||
import os
|
||||
import glob
|
||||
import typing
|
||||
|
||||
import imageio
|
||||
import numpy as np
|
||||
import json
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch.utils.data import Dataset
|
||||
|
||||
from body_model import MANO_JOINTS
|
||||
assert len(MANO_JOINTS) == 16
|
||||
import time
|
||||
|
||||
from util.logger import Logger
|
||||
from geometry.camera import invert_camera
|
||||
|
||||
from .tools import read_keypoints, read_mask_path, load_mano_preds, load_keypoints_with_interp
|
||||
from .vidproc import preprocess_cameras, preprocess_frames, preprocess_tracks
|
||||
|
||||
|
||||
"""
|
||||
Define data-related constants
|
||||
"""
|
||||
# DEFAULT_GROUND = np.array([0.0, -1.0, 0.0, -0.5])
|
||||
|
||||
# XXX: TEMPORARY CONSTANTS
|
||||
SHOT_PAD = 0
|
||||
MIN_SEQ_LEN = 10
|
||||
MAX_NUM_TRACKS = 12
|
||||
MIN_TRACK_LEN = 60
|
||||
MIN_KEYP_CONF = 0.4
|
||||
|
||||
|
||||
def get_dataset_from_cfg(cfg):
|
||||
args = cfg.data
|
||||
if not args.use_cams:
|
||||
args.sources.cameras = ""
|
||||
|
||||
args.sources = expand_source_paths(args.sources)
|
||||
print("DATA SOURCES", args.sources)
|
||||
a = time.time()
|
||||
check_data_sources(args, cfg)
|
||||
print('end of check_data_sources, time: ', time.time() - a)
|
||||
|
||||
return MultiPeopleDataset(
|
||||
args.sources,
|
||||
args.seq,
|
||||
tid_spec=args.track_ids,
|
||||
shot_idx=args.shot_idx,
|
||||
start_idx=int(args.start_idx),
|
||||
end_idx=int(args.end_idx),
|
||||
is_static=cfg.is_static,
|
||||
split_cameras=args.get("split_cameras", True),
|
||||
)
|
||||
|
||||
|
||||
def expand_source_paths(data_sources):
|
||||
return {k: get_data_source(v) for k, v in data_sources.items()}
|
||||
|
||||
|
||||
def get_data_source(source):
|
||||
matches = glob.glob(source)
|
||||
if len(matches) < 1:
|
||||
print(f"{source} does not exist")
|
||||
return source # return anyway for default values
|
||||
if len(matches) > 1:
|
||||
raise ValueError(f"{source} is not unique")
|
||||
return matches[0]
|
||||
|
||||
|
||||
def check_data_sources(args, cfg):
|
||||
print()
|
||||
a = time.time()
|
||||
if args.type == "video":
|
||||
preprocess_frames(args.sources.images, args.src_path, **args.frame_opts)
|
||||
b = time.time()
|
||||
preprocess_tracks(cfg.datatype, args.sources.images, args.sources.tracks, args.sources.shots, gpu=cfg.gpu)
|
||||
c = time.time()
|
||||
preprocess_cameras(args, overwrite=args.get("overwrite_cams", False))
|
||||
d = time.time()
|
||||
print('frame, hand track, slam camera: ', b-a, c-b, d-c)
|
||||
|
||||
|
||||
class MultiPeopleDataset(Dataset):
|
||||
def __init__(
|
||||
self,
|
||||
data_sources: typing.Dict,
|
||||
seq_name,
|
||||
tid_spec="all",
|
||||
shot_idx=0,
|
||||
start_idx=0,
|
||||
end_idx=-1,
|
||||
is_static=False,
|
||||
pad_shot=False,
|
||||
split_cameras=True,
|
||||
):
|
||||
self.is_static = is_static
|
||||
self.seq_name = seq_name
|
||||
self.data_sources = data_sources
|
||||
self.split_cameras = split_cameras
|
||||
|
||||
# select only images in the desired shot
|
||||
img_files, _ = get_shot_img_files(
|
||||
self.data_sources["shots"], shot_idx, pad_shot
|
||||
)
|
||||
end_idx = end_idx if end_idx > 0 else len(img_files)
|
||||
self.data_start, self.data_end = start_idx, end_idx
|
||||
img_files = img_files[start_idx:end_idx]
|
||||
self.img_names = [get_name(f) for f in img_files]
|
||||
self.num_imgs = len(self.img_names)
|
||||
|
||||
img_dir = self.data_sources["images"]
|
||||
assert os.path.isdir(img_dir)
|
||||
self.img_paths = [os.path.join(img_dir, f) for f in img_files]
|
||||
img_h, img_w = imageio.imread(self.img_paths[0]).shape[:2]
|
||||
self.img_size = img_w, img_h
|
||||
print(f"USING TOTAL {self.num_imgs} {img_w}x{img_h} IMGS")
|
||||
|
||||
# find the tracks in the video
|
||||
track_root = self.data_sources["tracks"]
|
||||
if tid_spec == "all" or tid_spec.startswith("longest"):
|
||||
n_tracks = MAX_NUM_TRACKS
|
||||
if tid_spec.startswith("longest"):
|
||||
n_tracks = int(tid_spec.split("-")[1])
|
||||
# get the longest tracks in the selected shot
|
||||
track_ids = sorted(os.listdir(track_root))
|
||||
track_paths = [
|
||||
[f"{track_root}/{tid}/{name}_keypoints.json" for name in self.img_names]
|
||||
for tid in track_ids
|
||||
]
|
||||
track_lens = [
|
||||
len(list(filter(os.path.isfile, paths))) for paths in track_paths
|
||||
]
|
||||
print("raw TRACK IDS and LENGTHS", track_ids, track_lens)
|
||||
track_ids = [
|
||||
track_ids[i]
|
||||
for i in range(len(track_lens))# np.argsort(track_lens)[::-1]
|
||||
if track_lens[i] > MIN_TRACK_LEN
|
||||
]
|
||||
print("TRACK LENGTHS", track_ids, track_lens)
|
||||
track_ids = track_ids[:n_tracks]
|
||||
else:
|
||||
track_ids = [f"{int(tid):03d}" for tid in tid_spec.split("-")]
|
||||
|
||||
print("TRACK IDS", track_ids)
|
||||
|
||||
self.track_ids = track_ids
|
||||
self.n_tracks = len(track_ids)
|
||||
self.track_dirs = [os.path.join(track_root, tid) for tid in track_ids]
|
||||
|
||||
# keep a list of frame index masks of whether a track is available in a frame
|
||||
sidx = np.inf
|
||||
eidx = -1
|
||||
self.track_vis_masks = []
|
||||
for pred_dir in self.track_dirs:
|
||||
kp_paths = [f"{pred_dir}/{x}_keypoints.json" for x in self.img_names]
|
||||
has_kp = [os.path.isfile(x) for x in kp_paths]
|
||||
|
||||
# keep track of which frames this track is visible in
|
||||
vis_mask = np.array(has_kp)
|
||||
idcs = np.where(vis_mask)[0]
|
||||
if len(idcs) > 0:
|
||||
si, ei = min(idcs), max(idcs)
|
||||
sidx = min(sidx, si)
|
||||
eidx = max(eidx, ei)
|
||||
self.track_vis_masks.append(vis_mask)
|
||||
|
||||
eidx = max(eidx + 1, 0)
|
||||
sidx = min(sidx, eidx)
|
||||
print("START", sidx, "END", eidx)
|
||||
self.start_idx = sidx
|
||||
self.end_idx = eidx
|
||||
self.seq_len = eidx - sidx
|
||||
self.seq_intervals = [(sidx, eidx) for _ in track_ids]
|
||||
|
||||
self.sel_img_paths = self.img_paths[sidx:eidx]
|
||||
self.sel_img_names = self.img_names[sidx:eidx]
|
||||
|
||||
# used to cache data
|
||||
self.data_dict = {}
|
||||
self.cam_data = None
|
||||
|
||||
def __len__(self):
|
||||
return self.n_tracks
|
||||
|
||||
def load_data(self, interp_input=True):
|
||||
if len(self.data_dict) > 0:
|
||||
return
|
||||
|
||||
# load camera data
|
||||
self.load_camera_data()
|
||||
# get data for each track
|
||||
data_out = {
|
||||
"mask_paths": [],
|
||||
"floor_plane": [],
|
||||
"joints2d": [],
|
||||
"vis_mask": [],
|
||||
"track_interval": [],
|
||||
"init_body_pose": [],
|
||||
"init_body_shape": [],
|
||||
"init_root_orient": [],
|
||||
"init_trans": [],
|
||||
"is_right": []
|
||||
}
|
||||
|
||||
# create batches of sequences
|
||||
# each batch is a track for a person
|
||||
T = self.seq_len
|
||||
sidx, eidx = self.start_idx, self.end_idx
|
||||
for i, tid in enumerate(self.track_ids):
|
||||
if len(self.track_ids) == 2:
|
||||
assert int(i) == int(tid), f'{i}/{tid}/{self.track_ids}'
|
||||
# load mask of visible frames for this track
|
||||
vis_mask = self.track_vis_masks[i][sidx:eidx] # (T)
|
||||
vis_idcs = np.where(vis_mask)[0]
|
||||
track_s, track_e = min(vis_idcs), max(vis_idcs) + 1
|
||||
data_out["track_interval"].append([track_s, track_e])
|
||||
|
||||
vis_mask = get_ternary_mask(vis_mask)
|
||||
data_out["vis_mask"].append(vis_mask)
|
||||
|
||||
# load 2d keypoints for visible frames with interpolation
|
||||
kp_paths = [
|
||||
f"{self.track_dirs[i]}/{x}_keypoints.json" for x in self.sel_img_names
|
||||
]
|
||||
# (T, J, 3) (x, y, conf) - with interpolation for missing frames
|
||||
joints2d_data = load_keypoints_with_interp(kp_paths, interp=interp_input)
|
||||
# print('joints2d_data: ', joints2d_data.shape)
|
||||
# print('joints2d_data[:, :, [2]]: ', np.unique(joints2d_data[:, :, [2]]))
|
||||
# raise ValueError
|
||||
# Discard bad ViTPose detections
|
||||
joints2d_data[:, :, 2] = 1.0
|
||||
joints2d_data[
|
||||
np.repeat(joints2d_data[:, :, [2]] < MIN_KEYP_CONF, 3, axis=2)
|
||||
] = 0
|
||||
# Set all confidence values to 1.0 for optimization
|
||||
data_out["joints2d"].append(joints2d_data)
|
||||
|
||||
# load single image mano predictions
|
||||
pred_paths = [
|
||||
f"{self.track_dirs[i]}/{x}_mano.json" for x in self.sel_img_names
|
||||
]
|
||||
pose_init, orient_init, trans_init, betas_init, is_right = load_mano_preds(
|
||||
pred_paths, tid=tid, interp=interp_input
|
||||
)
|
||||
|
||||
n_joints = len(MANO_JOINTS) - 1 # 15
|
||||
data_out["init_body_pose"].append(pose_init)
|
||||
data_out["init_body_shape"].append(betas_init)
|
||||
data_out["init_root_orient"].append(orient_init)
|
||||
data_out["init_trans"].append(trans_init)
|
||||
data_out['is_right'].append(is_right)
|
||||
|
||||
# DEBUG: Print hand type for each track
|
||||
print(f"DEBUG: Track {i} (tid={tid}): is_right={is_right[0] if len(is_right) > 0 else 'empty'}, {len(pred_paths)} frames")
|
||||
|
||||
# data_out["floor_plane"].append(DEFAULT_GROUND[:3] * DEFAULT_GROUND[3:])
|
||||
|
||||
self.data_dict = data_out
|
||||
|
||||
def __getitem__(self, idx):
|
||||
if len(self.data_dict) < 1:
|
||||
self.load_data()
|
||||
|
||||
obs_data = dict()
|
||||
|
||||
# 2D keypoints
|
||||
joint2d_data = self.data_dict["joints2d"][idx]
|
||||
obs_data["joints2d"] = torch.Tensor(joint2d_data)
|
||||
|
||||
# single frame predictions
|
||||
obs_data["init_body_pose"] = torch.Tensor(self.data_dict["init_body_pose"][idx])
|
||||
obs_data["init_body_shape"] = torch.Tensor(self.data_dict["init_body_shape"][idx])
|
||||
obs_data["init_root_orient"] = torch.Tensor(
|
||||
self.data_dict["init_root_orient"][idx]
|
||||
)
|
||||
obs_data["init_trans"] = torch.Tensor(self.data_dict["init_trans"][idx])
|
||||
obs_data["is_right"] = torch.Tensor(self.data_dict["is_right"][idx])
|
||||
|
||||
# floor plane
|
||||
# obs_data["floor_plane"] = torch.Tensor(self.data_dict["floor_plane"][idx])
|
||||
|
||||
# the frames the track is visible in
|
||||
obs_data["vis_mask"] = torch.Tensor(self.data_dict["vis_mask"][idx])
|
||||
|
||||
# the frames used in this subsequence
|
||||
obs_data["seq_interval"] = torch.Tensor(list(self.seq_intervals[idx])).to(
|
||||
torch.int
|
||||
)
|
||||
# the start and end interval of available keypoints
|
||||
obs_data["track_interval"] = torch.Tensor(
|
||||
self.data_dict["track_interval"][idx]
|
||||
).int()
|
||||
|
||||
print(len(self.track_ids), self.track_ids[idx], obs_data["is_right"].shape, obs_data["seq_interval"].shape)
|
||||
|
||||
obs_data["track_id"] = int(self.track_ids[idx])
|
||||
# print(self.track_ids[idx], obs_data["is_right"])
|
||||
# print(self.track_ids[idx] == obs_data["is_right"])
|
||||
# is_right should be constant across time for each track
|
||||
assert torch.all(obs_data["is_right"] == obs_data["is_right"][0]), "is_right should be constant"
|
||||
assert int(self.track_ids[idx]) == int(obs_data["is_right"][0]), f"Track ID {self.track_ids[idx]} != is_right {obs_data['is_right'][0]}"
|
||||
obs_data["seq_name"] = self.seq_name
|
||||
return obs_data
|
||||
|
||||
def load_camera_data(self):
|
||||
cam_dir = self.data_sources["cameras"]
|
||||
data_interval = 0, -1
|
||||
if self.split_cameras:
|
||||
data_interval = self.data_start, self.data_end
|
||||
track_interval = self.start_idx, self.end_idx
|
||||
self.cam_data = CameraData(
|
||||
cam_dir, self.seq_len, self.img_size, self.is_static, data_interval, track_interval
|
||||
)
|
||||
|
||||
def get_camera_data(self):
|
||||
if self.cam_data is None:
|
||||
raise ValueError
|
||||
return self.cam_data.as_dict()
|
||||
|
||||
|
||||
class CameraData(object):
|
||||
def __init__(
|
||||
self, cam_dir, seq_len, img_size, is_static, data_interval=[0, -1], track_interval=[0, -1]
|
||||
):
|
||||
self.img_size = img_size
|
||||
self.cam_dir = cam_dir
|
||||
self.is_static = is_static
|
||||
|
||||
# inclusive exclusive
|
||||
data_start, data_end = data_interval
|
||||
if data_end < 0:
|
||||
data_end += seq_len + 1
|
||||
data_len = data_end - data_start
|
||||
|
||||
# start and end indices are with respect to the data interval
|
||||
sidx, eidx = track_interval
|
||||
if eidx < 0:
|
||||
eidx += data_len + 1
|
||||
|
||||
self.sidx, self.eidx = sidx + data_start, eidx + data_start
|
||||
print(self.sidx, self.eidx, data_start, data_end, data_interval, track_interval)
|
||||
#
|
||||
self.seq_len = self.eidx - self.sidx
|
||||
|
||||
self.load_data()
|
||||
|
||||
def load_data(self):
|
||||
# camera info
|
||||
sidx, eidx = self.sidx, self.eidx
|
||||
img_w, img_h = self.img_size
|
||||
fpath = os.path.join(self.cam_dir, "cameras.npz")
|
||||
if os.path.isfile(fpath) and not self.is_static:
|
||||
Logger.log(f"Loading cameras from {fpath}...")
|
||||
cam_R, cam_t, intrins, width, height = load_cameras_npz(fpath, self.seq_len)
|
||||
scale = img_w / width
|
||||
self.intrins = scale * intrins[sidx:eidx]
|
||||
# move first camera to origin
|
||||
# R0, t0 = invert_camera(cam_R[sidx], cam_t[sidx])
|
||||
# self.cam_R = torch.einsum("ij,...jk->...ik", R0, cam_R[sidx:eidx])
|
||||
# self.cam_t = t0 + torch.einsum("ij,...j->...i", R0, cam_t[sidx:eidx])
|
||||
# t0 = -cam_t[sidx:eidx].mean(dim=0) + torch.randn(3) * 0.1
|
||||
t0 = -cam_t[sidx:sidx+1] + torch.randn(3) * 0.1
|
||||
self.cam_R = cam_R[sidx:eidx]
|
||||
self.cam_t = cam_t[sidx:eidx] - t0
|
||||
else:
|
||||
# raise ValueError
|
||||
Logger.log(f"WARNING: {fpath} does not exist, using static cameras...")
|
||||
default_focal = 0.5 * (img_h + img_w)
|
||||
self.intrins = torch.tensor(
|
||||
[default_focal, default_focal, img_w / 2, img_h / 2]
|
||||
)[None].repeat(self.seq_len, 1)
|
||||
|
||||
self.cam_R = torch.eye(3)[None].repeat(self.seq_len, 1, 1)
|
||||
self.cam_t = torch.zeros(self.seq_len, 3)
|
||||
|
||||
Logger.log(f"Images have {img_w}x{img_h}, intrins {self.intrins[0]}")
|
||||
print("CAMERA DATA", self.cam_R.shape, self.cam_t.shape, self.intrins[0])
|
||||
|
||||
def world2cam(self):
|
||||
return self.cam_R, self.cam_t
|
||||
|
||||
def cam2world(self):
|
||||
print('dataset.py, cam2world: ', self.cam_R)
|
||||
R = self.cam_R.transpose(-1, -2)
|
||||
t = -torch.einsum("bij,bj->bi", R, self.cam_t)
|
||||
print('dataset.py, cam2world: ', R)
|
||||
return R, t
|
||||
|
||||
def as_dict(self):
|
||||
return {
|
||||
"cam_R": self.cam_R, # (T, 3, 3)
|
||||
"cam_t": self.cam_t, # (T, 3)
|
||||
"intrins": self.intrins, # (T, 4)
|
||||
"static": self.is_static, # bool
|
||||
}
|
||||
|
||||
|
||||
def get_ternary_mask(vis_mask):
|
||||
# get the track start and end idcs relative to the filtered interval
|
||||
vis_mask = torch.as_tensor(vis_mask)
|
||||
vis_idcs = torch.where(vis_mask)[0]
|
||||
track_s, track_e = min(vis_idcs), max(vis_idcs) + 1
|
||||
# -1 = track out of scene, 0 = occlusion, 1 = visible
|
||||
vis_mask = vis_mask.float()
|
||||
vis_mask[:track_s] = -1
|
||||
vis_mask[track_e:] = -1
|
||||
return vis_mask
|
||||
|
||||
|
||||
def get_shot_img_files(shots_path, shot_idx, shot_pad=SHOT_PAD):
|
||||
assert os.path.isfile(shots_path)
|
||||
with open(shots_path, "r") as f:
|
||||
shots_dict = json.load(f)
|
||||
img_names = sorted(shots_dict.keys())
|
||||
N = len(img_names)
|
||||
shot_mask = np.array([shots_dict[x] == shot_idx for x in img_names])
|
||||
|
||||
idcs = np.where(shot_mask)[0]
|
||||
if shot_pad > 0: # drop the frames before/after shot change
|
||||
raise ValueError
|
||||
if min(idcs) > 0:
|
||||
idcs = idcs[shot_pad:]
|
||||
if len(idcs) > 0 and max(idcs) < N - 1:
|
||||
idcs = idcs[:-shot_pad]
|
||||
if len(idcs) < MIN_SEQ_LEN:
|
||||
raise ValueError("shot is too short for optimization")
|
||||
|
||||
shot_mask = np.zeros(N, dtype=bool)
|
||||
shot_mask[idcs] = 1
|
||||
|
||||
sel_paths = [img_names[i] for i in idcs]
|
||||
print(f"FOUND {len(idcs)}/{len(shots_dict)} FRAMES FOR SHOT {shot_idx}")
|
||||
return sel_paths, idcs
|
||||
|
||||
|
||||
def load_cameras_npz(camera_path, v):
|
||||
assert os.path.splitext(camera_path)[-1] == ".npz"
|
||||
|
||||
cam_data = np.load(camera_path)
|
||||
height, width, focal = (
|
||||
int(cam_data["height"]),
|
||||
int(cam_data["width"]),
|
||||
float(cam_data["focal"]),
|
||||
)
|
||||
|
||||
w2c = torch.from_numpy(cam_data["w2c"]) # (N, 4, 4)
|
||||
cam_R = w2c[:, :3, :3] # (N, 3, 3)
|
||||
cam_t = w2c[:, :3, 3] # (N, 3)
|
||||
N = len(w2c)
|
||||
|
||||
if "intrins" in cam_data:
|
||||
intrins = torch.from_numpy(cam_data["intrins"].astype(np.float32))
|
||||
else:
|
||||
intrins = torch.tensor([focal, focal, width / 2, height / 2])[None].repeat(N, 1)
|
||||
|
||||
print(f"Loaded {N} cameras")
|
||||
return cam_R, cam_t, intrins, width, height
|
||||
|
||||
|
||||
def is_image(x):
|
||||
return (x.endswith(".png") or x.endswith(".jpg")) and not x.startswith(".")
|
||||
|
||||
|
||||
def get_name(x):
|
||||
return os.path.splitext(os.path.basename(x))[0]
|
||||
|
||||
|
||||
def split_name(x, suffix):
|
||||
return os.path.basename(x).split(suffix)[0]
|
||||
|
||||
|
||||
def get_names_in_dir(d, suffix):
|
||||
files = [split_name(x, suffix) for x in glob.glob(f"{d}/*{suffix}")]
|
||||
return sorted(files)
|
||||
|
||||
|
||||
def batch_join(parent, names, suffix=""):
|
||||
return [os.path.join(parent, f"{n}{suffix}") for n in names]
|
||||
+168
@@ -0,0 +1,168 @@
|
||||
import os
|
||||
import json
|
||||
import functools
|
||||
|
||||
import numpy as np
|
||||
|
||||
from body_model import OP_NUM_JOINTS
|
||||
from scipy.interpolate import interp1d
|
||||
from scipy.spatial.transform import Rotation, Slerp
|
||||
|
||||
|
||||
def read_keypoints(keypoint_fn):
|
||||
"""
|
||||
Only reads body keypoint data of first person.
|
||||
"""
|
||||
empty_kps = np.zeros((OP_NUM_JOINTS, 3), dtype=np.float32)
|
||||
if not os.path.isfile(keypoint_fn):
|
||||
return empty_kps
|
||||
|
||||
with open(keypoint_fn) as keypoint_file:
|
||||
data = json.load(keypoint_file)
|
||||
|
||||
if len(data["people"]) == 0:
|
||||
print("WARNING: Found no keypoints in %s! Returning zeros!" % (keypoint_fn))
|
||||
return empty_kps
|
||||
|
||||
person_data = data["people"][0]
|
||||
body_keypoints = np.array(person_data["pose_keypoints_2d"], dtype=np.float32)
|
||||
body_keypoints = body_keypoints.reshape([-1, 3])
|
||||
return body_keypoints
|
||||
|
||||
|
||||
def load_keypoints_with_interp(kp_paths, interp=True):
|
||||
"""
|
||||
Load 2D keypoints from paths and interpolate missing frames between tmin and tmax.
|
||||
|
||||
Args:
|
||||
kp_paths: List of paths to keypoint JSON files
|
||||
interp: Whether to interpolate missing frames
|
||||
|
||||
Returns:
|
||||
joints2d_data: (T, J, 3) array of keypoints with interpolation
|
||||
"""
|
||||
# Load all keypoints
|
||||
joints2d_data = np.stack([read_keypoints(p) for p in kp_paths], axis=0).astype(np.float32)
|
||||
|
||||
if not interp:
|
||||
return joints2d_data
|
||||
|
||||
# Find which frames have valid keypoints (file exists and has data)
|
||||
vis_mask = np.array([os.path.isfile(p) for p in kp_paths])
|
||||
# Also check if keypoints are not all zeros
|
||||
has_data = ~np.all(joints2d_data == 0, axis=(1, 2))
|
||||
vis_mask = vis_mask & has_data
|
||||
vis_idcs = np.where(vis_mask)[0]
|
||||
|
||||
if len(vis_idcs) == 0 or len(vis_idcs) == 1:
|
||||
# No interpolation needed
|
||||
return joints2d_data
|
||||
|
||||
# Interpolate only between tmin and tmax (same as MANO interpolation)
|
||||
T, J, _ = joints2d_data.shape
|
||||
tmin, tmax = min(vis_idcs), max(vis_idcs) + 1
|
||||
times = np.arange(tmin, tmax)
|
||||
|
||||
# Interpolate each joint's x, y coordinates separately
|
||||
for j in range(J):
|
||||
# Interpolate x coordinate
|
||||
x_interp = interp1d(vis_idcs, joints2d_data[vis_idcs, j, 0],
|
||||
kind='linear', bounds_error=False)
|
||||
# Interpolate y coordinate
|
||||
y_interp = interp1d(vis_idcs, joints2d_data[vis_idcs, j, 1],
|
||||
kind='linear', bounds_error=False)
|
||||
|
||||
# Interpolate for frames between tmin and tmax
|
||||
for t in times:
|
||||
if t not in vis_idcs:
|
||||
# Interpolate
|
||||
joints2d_data[t, j, 0] = x_interp(t)
|
||||
joints2d_data[t, j, 1] = y_interp(t)
|
||||
|
||||
# Set confidence as average of neighboring visible frames
|
||||
prev_vis = vis_idcs[vis_idcs < t]
|
||||
next_vis = vis_idcs[vis_idcs > t]
|
||||
if len(prev_vis) > 0 and len(next_vis) > 0:
|
||||
conf_prev = joints2d_data[prev_vis[-1], j, 2]
|
||||
conf_next = joints2d_data[next_vis[0], j, 2]
|
||||
joints2d_data[t, j, 2] = (conf_prev + conf_next) / 2.0 * 0.8 # Slightly lower confidence for interpolated
|
||||
elif len(prev_vis) > 0:
|
||||
joints2d_data[t, j, 2] = joints2d_data[prev_vis[-1], j, 2] * 0.8
|
||||
elif len(next_vis) > 0:
|
||||
joints2d_data[t, j, 2] = joints2d_data[next_vis[0], j, 2] * 0.8
|
||||
|
||||
return joints2d_data
|
||||
|
||||
|
||||
def read_mask_path(path):
|
||||
mask_path = None
|
||||
if not os.path.isfile(path):
|
||||
return mask_path
|
||||
|
||||
with open(path, "r") as f:
|
||||
data = json.load(path)
|
||||
|
||||
person_data = data["people"][0]
|
||||
if "mask_path" in person_data:
|
||||
mask_path = person_data["mask_path"]
|
||||
|
||||
return mask_path
|
||||
|
||||
|
||||
def read_mano_preds(pred_path, tid, num_betas=10):
|
||||
"""
|
||||
reads the betas, body_pose, global orientation and translation of a mano prediction
|
||||
exported from phalp outputs
|
||||
returns betas (10,), body_pose (23, 3), global_orientation (3,), translation (3,)
|
||||
"""
|
||||
pose = np.zeros((15, 3))
|
||||
rot = np.zeros(3)
|
||||
trans = np.zeros(3)
|
||||
betas = np.zeros(num_betas)
|
||||
if not os.path.isfile(pred_path):
|
||||
return pose, rot, trans, betas, int(tid)
|
||||
|
||||
with open(pred_path, "r") as f:
|
||||
data = json.load(f)
|
||||
|
||||
pose = np.array(data["body_pose"], dtype=np.float32)
|
||||
rot = np.array(data["global_orient"], dtype=np.float32)
|
||||
trans = np.array(data["cam_trans"], dtype=np.float32)
|
||||
betas = np.array(data["betas"], dtype=np.float32)
|
||||
is_right = np.array(data["is_right"], dtype=np.float32)
|
||||
|
||||
return pose, rot, trans, betas, is_right
|
||||
|
||||
|
||||
def load_mano_preds(pred_paths, tid, interp=True, num_betas=10):
|
||||
vis_mask = np.array([os.path.isfile(x) for x in pred_paths])
|
||||
vis_idcs = np.where(vis_mask)[0]
|
||||
|
||||
# load single image mano predictions
|
||||
stack_fnc = functools.partial(np.stack, axis=0)
|
||||
# (N, 16, 3), (N, 3), (N, 3), (N, 10)
|
||||
pose, orient, trans, betas, is_right = map(
|
||||
stack_fnc, zip(*[read_mano_preds(p, tid, num_betas=num_betas) for p in pred_paths])
|
||||
)
|
||||
|
||||
assert len(np.where(is_right!=int(tid))[0]) == 0
|
||||
if not interp:
|
||||
return pose, orient, trans, betas, is_right
|
||||
|
||||
# interpolate the occluded tracks
|
||||
orient_slerp = Slerp(vis_idcs, Rotation.from_rotvec(orient[vis_idcs]))
|
||||
trans_interp = interp1d(vis_idcs, trans[vis_idcs], axis=0)
|
||||
betas_interp = interp1d(vis_idcs, betas[vis_idcs], axis=0)
|
||||
|
||||
tmin, tmax = min(vis_idcs), max(vis_idcs) + 1
|
||||
times = np.arange(tmin, tmax)
|
||||
orient[times] = orient_slerp(times).as_rotvec()
|
||||
trans[times] = trans_interp(times)
|
||||
betas[times] = betas_interp(times)
|
||||
|
||||
# interpolate for each joint angle
|
||||
for i in range(pose.shape[1]):
|
||||
pose_slerp = Slerp(vis_idcs, Rotation.from_rotvec(pose[vis_idcs, i]))
|
||||
pose[times, i] = pose_slerp(times).as_rotvec()
|
||||
|
||||
return pose, orient, trans, betas, is_right
|
||||
+284
@@ -0,0 +1,284 @@
|
||||
import os
|
||||
import shlex
|
||||
import sys
|
||||
from pathlib import Path
|
||||
import numpy as np
|
||||
import subprocess
|
||||
import cv2
|
||||
|
||||
import preproc.launch_hamer as hamer
|
||||
from preproc.launch_slam import split_frames_shots, get_command, check_intrins
|
||||
from preproc.extract_frames import video_to_frames, split_frame
|
||||
|
||||
|
||||
def is_nonempty(d):
|
||||
return os.path.isdir(d) and len(os.listdir(d)) > 0
|
||||
|
||||
|
||||
def preprocess_frames(img_dir, src_path, overwrite=False, **kwargs):
|
||||
if not overwrite and is_nonempty(img_dir):
|
||||
print(f"FOUND {len(os.listdir(img_dir))} FRAMES in {img_dir}")
|
||||
return
|
||||
print(f"EXTRACTING FRAMES FROM {src_path} TO {img_dir}")
|
||||
print(kwargs)
|
||||
|
||||
# out = video_to_frames(src_path, img_dir, overwrite=overwrite, **kwargs)
|
||||
out = split_frame(src_path, img_dir, overwrite=overwrite, **kwargs)
|
||||
assert out == 0, "FAILED FRAME EXTRACTION"
|
||||
|
||||
|
||||
def preprocess_tracks(datatype, img_dir, track_dir, shot_dir, gpu, overwrite=False):
|
||||
"""
|
||||
:param img_dir
|
||||
:param track_dir, expected format: res_root/track_name/sequence
|
||||
:param shot_dir, expected format: res_root/shot_name/sequence
|
||||
"""
|
||||
if not overwrite and is_nonempty(track_dir):
|
||||
print(f"FOUND TRACKS IN {track_dir}")
|
||||
return
|
||||
|
||||
print(f"RUNNING HAMER ON {img_dir}")
|
||||
track_root, seq = os.path.split(track_dir.rstrip("/"))
|
||||
res_root, track_name = os.path.split(track_root)
|
||||
shot_name = shot_dir.rstrip("/").split("/")[-2]
|
||||
|
||||
hamer.process_seq(
|
||||
[gpu],
|
||||
res_root,
|
||||
seq,
|
||||
img_dir,
|
||||
track_name=track_name,
|
||||
shot_name=shot_name,
|
||||
datatype=datatype,
|
||||
overwrite=overwrite,
|
||||
)
|
||||
|
||||
|
||||
def load_vipe_cameras(vipe_dir, seq_name, img_dir, start=0, end=-1):
|
||||
"""
|
||||
Load VIPE camera outputs and convert to DROID-SLAM format
|
||||
|
||||
Args:
|
||||
vipe_dir: Path to VIPE results directory
|
||||
seq_name: Sequence name
|
||||
img_dir: Image directory to get image size
|
||||
start: Start frame index
|
||||
end: End frame index
|
||||
|
||||
Returns:
|
||||
w2c: (N, 4, 4) world-to-camera matrices
|
||||
intrins_full: (N, 6) intrinsics [fx, fy, cx, cy, W, H]
|
||||
"""
|
||||
# Load VIPE pose and intrinsics
|
||||
vipe_pose_path = os.path.join(vipe_dir, "pose", f"{seq_name}.npz")
|
||||
vipe_intrins_path = os.path.join(vipe_dir, "intrinsics", f"{seq_name}.npz")
|
||||
|
||||
if not os.path.exists(vipe_pose_path):
|
||||
raise FileNotFoundError(f"VIPE pose file not found: {vipe_pose_path}")
|
||||
if not os.path.exists(vipe_intrins_path):
|
||||
raise FileNotFoundError(f"VIPE intrinsics file not found: {vipe_intrins_path}")
|
||||
|
||||
print(f"Loading VIPE cameras from {vipe_dir}")
|
||||
pose_data = np.load(vipe_pose_path)
|
||||
c2w = pose_data['data'] # (N, 4, 4) camera-to-world
|
||||
pose_inds = pose_data['inds'] # (N,) frame indices
|
||||
|
||||
intrins_data = np.load(vipe_intrins_path)
|
||||
intrins = intrins_data['data'] # (N, 4) [fx, fy, cx, cy]
|
||||
intrins_inds = intrins_data['inds'] # (N,) frame indices
|
||||
|
||||
# Verify indices match
|
||||
assert np.array_equal(pose_inds, intrins_inds), "Pose and intrinsics indices don't match!"
|
||||
|
||||
# Get image size from first image
|
||||
image_files = sorted([f for f in os.listdir(img_dir)
|
||||
if f.endswith(('.png', '.jpg', '.jpeg'))])
|
||||
if not image_files:
|
||||
# Infer from intrinsics (cx, cy should be roughly at center)
|
||||
img_width = int(intrins[0, 2] * 2)
|
||||
img_height = int(intrins[0, 3] * 2)
|
||||
print(f"Inferred image size from intrinsics: {img_width}x{img_height}")
|
||||
else:
|
||||
img_path = os.path.join(img_dir, image_files[0])
|
||||
img = cv2.imread(img_path)
|
||||
img_height, img_width = img.shape[:2]
|
||||
print(f"Got image size from images: {img_width}x{img_height}")
|
||||
|
||||
# Select frames based on start/end
|
||||
if end < 0:
|
||||
end = len(c2w)
|
||||
c2w = c2w[start:end]
|
||||
intrins = intrins[start:end]
|
||||
|
||||
# Convert camera-to-world to world-to-camera
|
||||
w2c = np.linalg.inv(c2w)
|
||||
|
||||
# Add width and height to intrinsics
|
||||
N = len(w2c)
|
||||
intrins_full = np.zeros((N, 6), dtype=np.float32)
|
||||
intrins_full[:, :4] = intrins
|
||||
intrins_full[:, 4] = img_width
|
||||
intrins_full[:, 5] = img_height
|
||||
|
||||
print(f"Loaded {N} VIPE camera poses")
|
||||
return w2c, intrins_full
|
||||
|
||||
|
||||
def save_vipe_cameras_as_droid(output_dir, w2c, intrins_full):
|
||||
"""
|
||||
Save VIPE cameras in DROID-SLAM format
|
||||
|
||||
Args:
|
||||
output_dir: Output directory
|
||||
w2c: (N, 4, 4) world-to-camera matrices
|
||||
intrins_full: (N, 6) intrinsics [fx, fy, cx, cy, W, H]
|
||||
"""
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
# Extract parameters
|
||||
W, H = intrins_full[0, 4], intrins_full[0, 5]
|
||||
focal = intrins_full[:, :2].mean()
|
||||
|
||||
print(f"Saving VIPE cameras to {output_dir}")
|
||||
print(f"Image size: {int(W)}x{int(H)}, focal: {focal:.2f}")
|
||||
|
||||
# Save cameras.npz (main format used by Dyn-HaMR)
|
||||
np.savez(
|
||||
f"{output_dir}/cameras.npz",
|
||||
height=H,
|
||||
width=W,
|
||||
focal=focal,
|
||||
intrins=intrins_full[:, :4],
|
||||
w2c=w2c,
|
||||
)
|
||||
print(f"Saved cameras.npz with {len(w2c)} frames")
|
||||
|
||||
|
||||
def run_vipe(video_path, vipe_dir, vipe_root=None, vipe_python=None, vipe_pipeline=None):
|
||||
"""Run VIPE without shell activation, using the project venv by default.
|
||||
|
||||
``VIPE_PYTHON`` or ``vipe_python`` can select a separate interpreter.
|
||||
Input/output paths are resolved before changing the child working directory.
|
||||
"""
|
||||
project_root = Path(__file__).resolve().parents[4]
|
||||
root = Path(vipe_root or project_root / "third_party" / "vipe").expanduser().resolve()
|
||||
video = Path(video_path).expanduser().absolute()
|
||||
output = Path(vipe_dir).expanduser().absolute()
|
||||
configured_python = vipe_python or os.environ.get("VIPE_PYTHON")
|
||||
project_python = project_root / ".venv" / "bin" / "python"
|
||||
python = Path(configured_python or (project_python if project_python.is_file() else sys.executable)).expanduser().absolute()
|
||||
if not video.is_file():
|
||||
raise FileNotFoundError(f"VIPE input video not found: {video}")
|
||||
if not (root / "vipe" / "cli" / "main.py").is_file():
|
||||
raise FileNotFoundError(f"VIPE source not found in {root}; set data.vipe_root")
|
||||
if not python.is_file():
|
||||
raise FileNotFoundError(f"VIPE interpreter not found: {python}; set VIPE_PYTHON")
|
||||
cmd = [str(python), "-m", "vipe.cli.main", "infer", str(video), "--output", str(output)]
|
||||
if vipe_pipeline:
|
||||
cmd.extend(["--pipeline", str(vipe_pipeline)])
|
||||
print(f"Running VIPE in {root}: {shlex.join(cmd)}", flush=True)
|
||||
out = subprocess.call(cmd, cwd=str(root))
|
||||
if out != 0:
|
||||
print(f"WARNING: VIPE failed with exit code {out}")
|
||||
else:
|
||||
print("VIPE completed successfully")
|
||||
return out
|
||||
|
||||
|
||||
def preprocess_cameras(cfg, overwrite=False):
|
||||
if not overwrite and is_nonempty(cfg.sources.cameras):
|
||||
print(f"FOUND CAMERAS IN {cfg.sources.cameras}")
|
||||
return
|
||||
|
||||
# Check if we should use VIPE instead of DROID-SLAM
|
||||
use_vipe = cfg.get("use_vipe", False)
|
||||
vipe_dir = cfg.get("vipe_dir", None)
|
||||
|
||||
if use_vipe and vipe_dir is not None:
|
||||
# Check if VIPE results exist for this sequence
|
||||
vipe_pose_path = os.path.join(vipe_dir, "pose", f"{cfg.seq}.npz")
|
||||
vipe_intrins_path = os.path.join(vipe_dir, "intrinsics", f"{cfg.seq}.npz")
|
||||
|
||||
# If VIPE results don't exist, run VIPE
|
||||
if not (os.path.exists(vipe_pose_path) and os.path.exists(vipe_intrins_path)):
|
||||
print(f"VIPE results not found for sequence '{cfg.seq}', running VIPE...")
|
||||
|
||||
# The source location is independent of the output directory.
|
||||
vipe_root = cfg.get("vipe_root", None)
|
||||
video_path = cfg.get("src_path", None)
|
||||
|
||||
if video_path is None or not os.path.exists(video_path):
|
||||
raise FileNotFoundError(
|
||||
f"Video path not found: {video_path}\n"
|
||||
f"Cannot run VIPE. Please provide a valid 'src_path' in your config."
|
||||
)
|
||||
|
||||
# Run VIPE
|
||||
out = run_vipe(video_path, vipe_dir, vipe_root, cfg.get("vipe_python", None), cfg.get("vipe_pipeline", None))
|
||||
if out != 0:
|
||||
raise RuntimeError(
|
||||
f"VIPE failed with exit code {out}\n"
|
||||
f"Check the command and error above. Select an installed Python "
|
||||
f"with VIPE_PYTHON or data.vipe_python; conda is not required."
|
||||
)
|
||||
|
||||
# Load VIPE results (after potentially running VIPE)
|
||||
if not (os.path.exists(vipe_pose_path) and os.path.exists(vipe_intrins_path)):
|
||||
raise FileNotFoundError(
|
||||
f"VIPE results not found after execution:\n"
|
||||
f" Pose: {vipe_pose_path}\n"
|
||||
f" Intrinsics: {vipe_intrins_path}\n"
|
||||
f"VIPE may have failed silently. Please check VIPE logs."
|
||||
)
|
||||
|
||||
print(f"USING VIPE CAMERAS FROM {vipe_dir}")
|
||||
img_dir = cfg.sources.images
|
||||
map_dir = cfg.sources.cameras
|
||||
|
||||
# Get frame range
|
||||
subseqs, shot_idcs = split_frames_shots(cfg.sources.images, cfg.sources.shots)
|
||||
shot_idx = np.where(shot_idcs == cfg.shot_idx)[0][0]
|
||||
start, end = subseqs[shot_idx]
|
||||
|
||||
if not cfg.split_cameras:
|
||||
# only run on specified segment within shot
|
||||
end = start + cfg.end_idx
|
||||
start = start + cfg.start_idx
|
||||
|
||||
# Load and convert VIPE cameras
|
||||
w2c, intrins_full = load_vipe_cameras(vipe_dir, cfg.seq, img_dir, start, end)
|
||||
save_vipe_cameras_as_droid(map_dir, w2c, intrins_full)
|
||||
return
|
||||
|
||||
# Default: use DROID-SLAM
|
||||
print(f"RUNNING SLAM ON {cfg.seq}")
|
||||
img_dir = cfg.sources.images
|
||||
map_dir = cfg.sources.cameras
|
||||
subseqs, shot_idcs = split_frames_shots(cfg.sources.images, cfg.sources.shots)
|
||||
print(shot_idcs, cfg.shot_idx, np.where(shot_idcs == cfg.shot_idx), cfg.sources.images, cfg.sources.shots)
|
||||
print(subseqs)
|
||||
shot_idx = np.where(shot_idcs == cfg.shot_idx)[0][0]
|
||||
# run on selected shot
|
||||
start, end = subseqs[shot_idx]
|
||||
if not cfg.split_cameras:
|
||||
# only run on specified segment within shot
|
||||
end = start + cfg.end_idx
|
||||
start = start + cfg.start_idx
|
||||
intrins_path = cfg.sources.get("intrins", None)
|
||||
if intrins_path is not None:
|
||||
intrins_path = check_intrins(cfg.type, cfg.root, intrins_path, cfg.seq, cfg.split)
|
||||
|
||||
print('img_dir, map_dir, start, end, intrins_path', img_dir, map_dir, start, end, intrins_path)
|
||||
# raise ValueERRROR
|
||||
cmd = get_command(
|
||||
img_dir,
|
||||
map_dir,
|
||||
start=start,
|
||||
end=end,
|
||||
intrins_path=intrins_path,
|
||||
overwrite=overwrite,
|
||||
)
|
||||
print(cmd)
|
||||
gpu = os.environ.get("CUDA_VISIBLE_DEVICES", 0)
|
||||
out = subprocess.call(f"CUDA_VISIBLE_DEVICES={gpu} {cmd}", shell=True)
|
||||
assert out == 0, "SLAM FAILED"
|
||||
@@ -0,0 +1,5 @@
|
||||
from . import camera
|
||||
from . import mesh
|
||||
from . import pcl
|
||||
from . import plane
|
||||
from . import rotation
|
||||
+351
@@ -0,0 +1,351 @@
|
||||
import torch
|
||||
import numpy as np
|
||||
|
||||
|
||||
def perspective_projection(
|
||||
points, focal_length, camera_center, rotation=None, translation=None
|
||||
):
|
||||
"""
|
||||
Adapted from https://github.com/mkocabas/VIBE/blob/master/lib/models/spin.py
|
||||
This function computes the perspective projection of a set of points.
|
||||
Input:
|
||||
points (bs, N, 3): 3D points
|
||||
focal_length (bs, 2): Focal length
|
||||
camera_center (bs, 2): Camera center
|
||||
rotation (bs, 3, 3): OPTIONAL Camera rotation
|
||||
translation (bs, 3): OPTIONAL Camera translation
|
||||
"""
|
||||
batch_size = points.shape[0]
|
||||
K = torch.zeros([batch_size, 3, 3], device=points.device)
|
||||
K[:, 0, 0] = focal_length[:, 0]
|
||||
K[:, 1, 1] = focal_length[:, 1]
|
||||
K[:, 2, 2] = 1.0
|
||||
K[:, :-1, -1] = camera_center
|
||||
|
||||
if rotation is not None and translation is not None:
|
||||
# Transform points
|
||||
points = torch.einsum("bij,bkj->bki", rotation, points)
|
||||
points = points + translation.unsqueeze(1)
|
||||
|
||||
# Apply perspective distortion
|
||||
projected_points = points / points[..., 2:3]
|
||||
|
||||
# Apply camera intrinsics
|
||||
projected_points = torch.einsum("bij,bkj->bki", K, projected_points)
|
||||
|
||||
return projected_points[:, :, :-1]
|
||||
|
||||
|
||||
def reproject(points3d, cam_R, cam_t, cam_f, cam_center):
|
||||
"""
|
||||
reproject points3d into the scene cameras
|
||||
:param points3d (B, T, N, 3)
|
||||
:param cam_R (B, T, 3, 3)
|
||||
:param cam_t (B, T, 3)
|
||||
:param cam_f (T, 2)
|
||||
:param cam_center (T, 2)
|
||||
"""
|
||||
# print(points3d.shape, cam_R.shape, cam_t.shape, cam_f.shape, cam_center.shape)
|
||||
# print(cam_f[None, :, None].shape, cam_center[None, :, None].shape)
|
||||
# raise ValueError
|
||||
B, T, N, _ = points3d.shape
|
||||
points3d = torch.einsum("btij,btnj->btni", cam_R, points3d)
|
||||
points3d = points3d + cam_t[..., None, :] # (B, T, N, 3)
|
||||
points2d = points3d[..., :2] / points3d[..., 2:3]
|
||||
points2d = cam_f[None, :, None] * points2d + cam_center[None, :, None]
|
||||
return points2d
|
||||
|
||||
|
||||
def focal2fov(focal, R):
|
||||
"""
|
||||
:param focal, focal length
|
||||
:param R, either W / 2 or H / 2
|
||||
"""
|
||||
return 2 * np.arctan(R / focal)
|
||||
|
||||
|
||||
def fov2focal(fov, R):
|
||||
"""
|
||||
:param fov, field of view in radians
|
||||
:param R, either W / 2 or H / 2
|
||||
"""
|
||||
return R / np.tan(fov / 2)
|
||||
|
||||
|
||||
def compute_lookat_box(bb_min, bb_max, intrins):
|
||||
"""
|
||||
The center and distance to a scene with bb_min, bb_max
|
||||
to place a camera with given intrinsics
|
||||
:param bb_min (3,)
|
||||
:param bb_max (3,)
|
||||
:param intrinsics, (fx, fy, cx, cy) of camera
|
||||
:param view_angle (optional) viewing angle in radians (elevation)
|
||||
"""
|
||||
fx, fy, cx, cy = intrins
|
||||
bb_min, bb_max = torch.tensor(bb_min), torch.tensor(bb_max)
|
||||
center = 0.5 * (bb_min + bb_max)
|
||||
size = torch.linalg.norm(bb_max - bb_min)
|
||||
cam_dist = np.sqrt(fx**2 + fy**2) / np.sqrt(cx**2 + cy**2)
|
||||
cam_dist = 0.75 * size * cam_dist
|
||||
return center, cam_dist
|
||||
|
||||
|
||||
def lookat_origin(cam_dist, view_angle=-np.pi / 6):
|
||||
"""
|
||||
:param cam_dist (float)
|
||||
:param view_angle (float)
|
||||
"""
|
||||
cam_dist = np.abs(cam_dist)
|
||||
view_angle = np.abs(view_angle)
|
||||
pos = cam_dist * torch.tensor([0, np.sin(view_angle), np.cos(view_angle)])
|
||||
rot = rotx(view_angle)
|
||||
return rot, pos
|
||||
|
||||
|
||||
def lookat_matrix(source_pos, target_pos, up):
|
||||
"""
|
||||
IMPORTANT: USES RIGHT UP BACK XYZ CONVENTION
|
||||
:param source_pos (*, 3)
|
||||
:param target_pos (*, 3)
|
||||
:param up (3,)
|
||||
"""
|
||||
*dims, _ = source_pos.shape
|
||||
up = up.reshape(*(1,) * len(dims), 3)
|
||||
up = up / torch.linalg.norm(up, dim=-1, keepdim=True)
|
||||
back = normalize(target_pos - source_pos)
|
||||
right = normalize(torch.linalg.cross(up, back))
|
||||
up = normalize(torch.linalg.cross(back, right))
|
||||
R = torch.stack([right, up, back], dim=-1)
|
||||
return make_4x4_pose(R, source_pos)
|
||||
|
||||
|
||||
def normalize(x):
|
||||
return x / torch.linalg.norm(x, dim=-1, keepdim=True)
|
||||
|
||||
|
||||
def invert_camera(R, t):
|
||||
"""
|
||||
:param R (*, 3, 3)
|
||||
:param t (*, 3)
|
||||
returns Ri (*, 3, 3), ti (*, 3)
|
||||
"""
|
||||
R, t = torch.tensor(R), torch.tensor(t)
|
||||
Ri = R.transpose(-1, -2)
|
||||
ti = -torch.einsum("...ij,...j->...i", Ri, t)
|
||||
return Ri, ti
|
||||
|
||||
|
||||
def compose_cameras(R1, t1, R2, t2):
|
||||
"""
|
||||
composes [R1, t1] and [R2, t2]
|
||||
:param R1 (*, 3, 3)
|
||||
:param t1 (*, 3)
|
||||
:param R2 (*, 3, 3)
|
||||
:param t2 (*, 3)
|
||||
"""
|
||||
R = torch.einsum("...ij,...jk->...ik", R1, R2)
|
||||
t = t1 + torch.einsum("...ij,...j->...i", R1, t2)
|
||||
return R, t
|
||||
|
||||
|
||||
def matmul_nd(A, x):
|
||||
"""
|
||||
multiply batch matrix A to batch nd tensors
|
||||
:param A (B, m, n)
|
||||
:param x (B, *dims, m)
|
||||
"""
|
||||
B, m, n = A.shape
|
||||
assert len(A) == len(x)
|
||||
assert x.shape[-1] == m
|
||||
B, *dims, _ = x.shape
|
||||
return torch.matmul(A.reshape(B, *(1,) * len(dims), m, n), x[..., None])[..., 0]
|
||||
|
||||
|
||||
def view_matrix(z, up, pos):
|
||||
"""
|
||||
:param z (*, 3) up (*, 3) pos (*, 3)
|
||||
returns (*, 4, 4)
|
||||
"""
|
||||
*dims, _ = z.shape
|
||||
x = normalize(torch.linalg.cross(up, z))
|
||||
y = normalize(torch.linalg.cross(z, x))
|
||||
bottom = (
|
||||
torch.tensor([0, 0, 0, 1], dtype=torch.float32)
|
||||
.reshape(*(1,) * len(dims), 1, 4)
|
||||
.expand(*dims, 1, 4)
|
||||
)
|
||||
|
||||
return torch.cat([torch.stack([x, y, z, pos], dim=-1), bottom], dim=-2)
|
||||
|
||||
|
||||
def average_pose(poses):
|
||||
"""
|
||||
:param poses (N, 4, 4)
|
||||
returns average pose (4, 4)
|
||||
"""
|
||||
center = poses[:, :3, 3].mean(0)
|
||||
up = normalize(poses[:, :3, 1].sum(0))
|
||||
z = normalize(poses[:, :3, 2].sum(0))
|
||||
return view_matrix(z, up, center)
|
||||
|
||||
|
||||
def project_so3(M, eps=1e-4):
|
||||
"""
|
||||
:param M (N, *, 3, 3)
|
||||
"""
|
||||
N, *dims, _, _ = M.shape
|
||||
M = M * (1 + torch.rand(N, *dims, 1, 3, device=M.device))
|
||||
U, D, Vt = torch.linalg.svd(M) # (N, *, 3, 3), (N, *, 3), (N, *, 3, 3)
|
||||
detuvt = torch.linalg.det(torch.matmul(U, Vt)) # (N, *)
|
||||
S = torch.cat(
|
||||
[torch.ones(N, *dims, 2, device=M.device), detuvt[..., None]], dim=-1
|
||||
) # (N, *, 3)
|
||||
return torch.matmul(U, torch.matmul(torch.diag_embed(S), Vt))
|
||||
|
||||
|
||||
def make_translation(t):
|
||||
return make_4x4_pose(torch.eye(3), t)
|
||||
|
||||
|
||||
def make_rotation(rx=0, ry=0, rz=0, order="xyz"):
|
||||
Rx = rotx(rx)
|
||||
Ry = roty(ry)
|
||||
Rz = rotz(rz)
|
||||
if order == "xyz":
|
||||
R = Rz @ Ry @ Rx
|
||||
elif order == "xzy":
|
||||
R = Ry @ Rz @ Rx
|
||||
elif order == "yxz":
|
||||
R = Rz @ Rx @ Ry
|
||||
elif order == "yzx":
|
||||
R = Rx @ Rz @ Ry
|
||||
elif order == "zyx":
|
||||
R = Rx @ Ry @ Rz
|
||||
elif order == "zxy":
|
||||
R = Ry @ Rx @ Rz
|
||||
return make_4x4_pose(R, torch.zeros(3))
|
||||
|
||||
|
||||
def make_4x4_pose(R, t):
|
||||
"""
|
||||
:param R (*, 3, 3)
|
||||
:param t (*, 3)
|
||||
return (*, 4, 4)
|
||||
"""
|
||||
dims = R.shape[:-2]
|
||||
pose_3x4 = torch.cat([R, t.view(*dims, 3, 1)], dim=-1)
|
||||
bottom = (
|
||||
torch.tensor([0, 0, 0, 1], device=R.device)
|
||||
.reshape(*(1,) * len(dims), 1, 4)
|
||||
.expand(*dims, 1, 4)
|
||||
)
|
||||
return torch.cat([pose_3x4, bottom], dim=-2)
|
||||
|
||||
|
||||
def normalize(x):
|
||||
return x / torch.sqrt(torch.sum(x**2, dim=-1, keepdim=True))
|
||||
|
||||
|
||||
def rotx(theta):
|
||||
return torch.tensor(
|
||||
[
|
||||
[1, 0, 0],
|
||||
[0, np.cos(theta), -np.sin(theta)],
|
||||
[0, np.sin(theta), np.cos(theta)],
|
||||
],
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
|
||||
def roty(theta):
|
||||
return torch.tensor(
|
||||
[
|
||||
[np.cos(theta), 0, np.sin(theta)],
|
||||
[0, 1, 0],
|
||||
[-np.sin(theta), 0, np.cos(theta)],
|
||||
],
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
|
||||
def rotz(theta):
|
||||
return torch.tensor(
|
||||
[
|
||||
[np.cos(theta), -np.sin(theta), 0],
|
||||
[np.sin(theta), np.cos(theta), 0],
|
||||
[0, 0, 1],
|
||||
],
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
|
||||
def relative_pose_c2w(Rwc1, Rwc2, twc1, twc2):
|
||||
"""
|
||||
compute relative pose from cam 1 to cam 2 given c2w pose matrices
|
||||
:param Rwc1, Rwc2 (N, 3, 3) cam1, cam2 to world rotations
|
||||
:param twc1, twc2 (N, 3) cam1, cam2 to world translations
|
||||
returns R21 (N, 3, 3) t21 (N, 3)
|
||||
"""
|
||||
twc1 = twc1.view(-1, 3, 1)
|
||||
twc2 = twc2.view(-1, 3, 1)
|
||||
Rc2w = Rwc2.transpose(-1, -2) # world to c2
|
||||
tc2w = -torch.matmul(Rc2w, twc2)
|
||||
Rc2c1 = torch.matmul(Rc2w, Rwc1)
|
||||
tc2c1 = tc2w + torch.matmul(Rc2w, twc1)
|
||||
return Rc2c1, tc2c1[..., 0]
|
||||
|
||||
|
||||
def relative_pose_w2c(Rc1w, Rc2w, tc1w, tc2w):
|
||||
"""
|
||||
compute relative pose from cam 1 to cam 2 given w2c camera matrices
|
||||
:param Rc1w, Rc2w (N, 3, 3) world to cam1, cam2 rotations
|
||||
:param tc1w, tc2w (N, 3) world to cam1, cam2 translations
|
||||
"""
|
||||
tc1w = tc1w.view(-1, 3, 1)
|
||||
tc2w = tc2w.view(-1, 3, 1)
|
||||
# we keep the world to cam transforms
|
||||
Rwc1 = Rc1w.transpose(-1, -2) # c1 to world
|
||||
twc1 = -torch.matmul(Rwc1, tc1w)
|
||||
Rc2c1 = torch.matmul(Rc2w, Rwc1) # c1 to c2
|
||||
tc2c1 = tc2w + torch.matmul(Rc2w, twc1)
|
||||
return Rc2c1, tc2c1[..., 0]
|
||||
|
||||
|
||||
def project(xyz_c, center, focal, eps=1e-5):
|
||||
"""
|
||||
:param xyz_c (*, 3) 3d point in camera coordinates
|
||||
:param focal (1)
|
||||
:param center (*, 2)
|
||||
return (*, 2)
|
||||
"""
|
||||
return focal * xyz_c[..., :2] / (xyz_c[..., 2:3] + eps) + center # (N, *, 2)
|
||||
|
||||
|
||||
def convert_yup(xyz):
|
||||
"""
|
||||
converts points in x right y down z forward to x right y up z back
|
||||
:param xyz (*, 3)
|
||||
"""
|
||||
x, y, z = torch.split(xyz[..., :3], 1, dim=-1)
|
||||
return torch.cat([x, -y, -z], dim=-1)
|
||||
|
||||
|
||||
def inv_project(uv, z, center, focal, yup=True):
|
||||
"""
|
||||
:param uv (*, 2)
|
||||
:param z (*, 1)
|
||||
:param center (*, 2)
|
||||
:param focal (1)
|
||||
:returns (*, 3)
|
||||
"""
|
||||
uv = uv - center
|
||||
if yup:
|
||||
return z * torch.cat(
|
||||
[uv[..., :1] / focal, -uv[..., 1:2] / focal, -torch.ones_like(uv[..., :1])],
|
||||
dim=-1,
|
||||
) # (N, *, 3)
|
||||
|
||||
return z * torch.cat(
|
||||
[uv / focal, torch.ones_like(uv[..., :1])],
|
||||
dim=-1,
|
||||
) # (N, *, 3)
|
||||
+145
@@ -0,0 +1,145 @@
|
||||
import os
|
||||
import numpy as np
|
||||
import trimesh
|
||||
|
||||
|
||||
def get_mesh_bb(mesh):
|
||||
"""
|
||||
:param mesh - trimesh mesh object
|
||||
returns bb_min (3), bb_max (3)
|
||||
"""
|
||||
bb_min = mesh.vertices.max(axis=0)
|
||||
bb_max = mesh.vertices.min(axis=0)
|
||||
return bb_min, bb_max
|
||||
|
||||
|
||||
def get_scene_bb(meshes):
|
||||
"""
|
||||
:param mesh_seqs - (potentially nested) list of trimesh objects
|
||||
returns bb_min (3), bb_max (3)
|
||||
"""
|
||||
if isinstance(meshes, trimesh.Trimesh):
|
||||
return get_mesh_bb(meshes)
|
||||
|
||||
bb_mins, bb_maxs = zip(*[get_scene_bb(mesh) for mesh in meshes])
|
||||
bb_mins = np.stack(bb_mins, axis=0)
|
||||
bb_maxs = np.stack(bb_maxs, axis=0)
|
||||
return bb_mins.min(axis=0), bb_maxs.max(axis=0)
|
||||
|
||||
|
||||
def make_batch_mesh(verts, faces, colors):
|
||||
"""
|
||||
convenience function to make batch of meshes
|
||||
meshs have same faces in batch, verts have same color in mesh
|
||||
:param verts (B, V, 3)
|
||||
:param faces (F, 3)
|
||||
:param colors (B, 3)
|
||||
"""
|
||||
B, V, _ = verts.shape
|
||||
return [make_mesh(verts[b], faces, colors[b, None].expand(V, -1)) for b in range(B)]
|
||||
|
||||
|
||||
def make_mesh(verts, faces, colors=None, yup=True):
|
||||
"""
|
||||
create a trimesh object for the faces and vertices
|
||||
:param verts (V, 3) tensor
|
||||
:param faces (F, 3) tensor
|
||||
:param colors (optional) (V, 3) tensor
|
||||
:param yup (optional bool) whether or not to save with Y up
|
||||
"""
|
||||
verts = verts.detach().cpu().numpy()
|
||||
faces = faces.detach().cpu().numpy()
|
||||
if yup:
|
||||
verts = np.array([1, -1, -1])[None, :] * verts
|
||||
if colors is None:
|
||||
colors = np.ones_like(verts) * 0.5
|
||||
else:
|
||||
colors = colors.detach().cpu().numpy()
|
||||
return trimesh.Trimesh(
|
||||
vertices=verts, faces=faces, vertex_colors=colors, process=False
|
||||
)
|
||||
|
||||
|
||||
def save_mesh_scenes(out_dir, scenes):
|
||||
"""
|
||||
:param scenes, list of scenes (list of meshes)
|
||||
"""
|
||||
assert isinstance(scenes, list)
|
||||
assert isinstance(scenes[0], list)
|
||||
B = len(scenes[0])
|
||||
if B == 1:
|
||||
save_meshes_to_obj(out_dir, [x[0] for x in scenes])
|
||||
else:
|
||||
save_scenes_to_glb(out_dir, scenes)
|
||||
|
||||
|
||||
def save_scenes_to_glb(out_dir, scenes):
|
||||
"""
|
||||
Saves a list of scenes (list of meshes) each to glb files
|
||||
"""
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
for t, meshes in enumerate(scenes):
|
||||
save_meshes_to_glb(f"{out_dir}/scene_{t:03d}.glb", meshes)
|
||||
|
||||
|
||||
def save_meshes_to_glb(path, meshes, names=None):
|
||||
"""
|
||||
put trimesh meshes in a scene and export to glb
|
||||
"""
|
||||
if names is not None:
|
||||
assert len(meshes) == len(names)
|
||||
|
||||
scene = trimesh.Scene()
|
||||
for i, mesh in enumerate(meshes):
|
||||
name = f"mesh_{i:03d}" if names is None else names[i]
|
||||
scene.add_geometry(mesh, node_name=name)
|
||||
|
||||
with open(path, "wb") as f:
|
||||
f.write(trimesh.exchange.gltf.export_glb(scene, include_normals=True))
|
||||
|
||||
|
||||
def save_meshes_to_obj(out_dir, meshes, names=None):
|
||||
if names is not None:
|
||||
assert len(meshes) == len(names)
|
||||
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
for i, mesh in enumerate(meshes):
|
||||
name = f"mesh_{i:03d}" if names is None else names[i]
|
||||
path = os.path.join(out_dir, f"{name}.obj")
|
||||
with open(path, "w") as f:
|
||||
mesh.export(f, file_type="obj", include_color=False, include_normals=True)
|
||||
|
||||
def vertices_to_trimesh(vertices, faces, mesh_base_color=(1.0, 1.0, 0.9),
|
||||
rot_axis=[1,0,0], rot_angle=0, is_right=1):
|
||||
# material = pyrender.MetallicRoughnessMaterial(
|
||||
# metallicFactor=0.0,
|
||||
# alphaMode='OPAQUE',
|
||||
# baseColorFactor=(*mesh_base_color, 1.0))
|
||||
faces_new = np.array([[92, 38, 234],
|
||||
[234, 38, 239],
|
||||
[38, 122, 239],
|
||||
[239, 122, 279],
|
||||
[122, 118, 279],
|
||||
[279, 118, 215],
|
||||
[118, 117, 215],
|
||||
[215, 117, 214],
|
||||
[117, 119, 214],
|
||||
[214, 119, 121],
|
||||
[119, 120, 121],
|
||||
[121, 120, 78],
|
||||
[120, 108, 78],
|
||||
[78, 108, 79]])
|
||||
faces = np.concatenate([faces, faces_new], axis=0)
|
||||
vertex_colors = np.array([(*mesh_base_color, 1.0)] * vertices.shape[0])
|
||||
# print(vertices.shape, self.faces.shape)
|
||||
mesh = trimesh.Trimesh(vertices.copy(), faces.copy(), vertex_colors=vertex_colors)
|
||||
# mesh = trimesh.Trimesh(vertices.copy(), self.faces.copy())
|
||||
|
||||
rot = trimesh.transformations.rotation_matrix(
|
||||
np.radians(rot_angle), rot_axis)
|
||||
mesh.apply_transform(rot)
|
||||
|
||||
rot = trimesh.transformations.rotation_matrix(
|
||||
np.radians(180), [1, 0, 0])
|
||||
mesh.apply_transform(rot)
|
||||
return mesh
|
||||
+60
@@ -0,0 +1,60 @@
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
|
||||
def read_pcl_tensor(path):
|
||||
pcl_np = read_pcl(path)
|
||||
return torch.from_numpy(pcl_np)
|
||||
|
||||
|
||||
def align_pcl(Y, X, weight=None, fixed_scale=False):
|
||||
"""align similarity transform to align X with Y using umeyama method
|
||||
X' = s * R * X + t is aligned with Y
|
||||
:param Y (*, N, 3) first trajectory
|
||||
:param X (*, N, 3) second trajectory
|
||||
:param weight (*, N, 1) optional weight of valid correspondences
|
||||
:returns s (*, 1), R (*, 3, 3), t (*, 3)
|
||||
"""
|
||||
*dims, N, _ = Y.shape
|
||||
N = torch.ones(*dims, 1, 1) * N
|
||||
|
||||
if weight is not None:
|
||||
Y = Y * weight
|
||||
X = X * weight
|
||||
N = weight.sum(dim=-2, keepdim=True) # (*, 1, 1)
|
||||
|
||||
# subtract mean
|
||||
my = Y.sum(dim=-2) / N[..., 0] # (*, 3)
|
||||
mx = X.sum(dim=-2) / N[..., 0]
|
||||
y0 = Y - my[..., None, :] # (*, N, 3)
|
||||
x0 = X - mx[..., None, :]
|
||||
|
||||
if weight is not None:
|
||||
y0 = y0 * weight
|
||||
x0 = x0 * weight
|
||||
|
||||
# correlation
|
||||
C = torch.matmul(y0.transpose(-1, -2), x0) / N # (*, 3, 3)
|
||||
U, D, Vh = torch.linalg.svd(C) # (*, 3, 3), (*, 3), (*, 3, 3)
|
||||
|
||||
S = torch.eye(3).reshape(*(1,) * (len(dims)), 3, 3).repeat(*dims, 1, 1)
|
||||
neg = torch.det(U) * torch.det(Vh.transpose(-1, -2)) < 0
|
||||
S[neg, 2, 2] = -1
|
||||
|
||||
R = torch.matmul(U, torch.matmul(S, Vh)) # (*, 3, 3)
|
||||
|
||||
D = torch.diag_embed(D) # (*, 3, 3)
|
||||
if fixed_scale:
|
||||
s = torch.ones(*dims, 1, device=Y.device, dtype=torch.float32)
|
||||
else:
|
||||
var = torch.sum(torch.square(x0), dim=(-1, -2), keepdim=True) / N # (*, 1, 1)
|
||||
s = (
|
||||
torch.diagonal(torch.matmul(D, S), dim1=-2, dim2=-1).sum(
|
||||
dim=-1, keepdim=True
|
||||
)
|
||||
/ var[..., 0]
|
||||
) # (*, 1)
|
||||
|
||||
t = my - s * torch.matmul(R, mx[..., None])[..., 0] # (*, 3)
|
||||
|
||||
return s, R, t
|
||||
+101
@@ -0,0 +1,101 @@
|
||||
import torch
|
||||
|
||||
|
||||
def fit_plane(points):
|
||||
"""
|
||||
:param points (*, N, 3)
|
||||
returns (*, 3) plane parameters (returns in normal * offset format)
|
||||
"""
|
||||
*dims, N, D = points.shape
|
||||
mean = points.mean(dim=-2, keepdim=True)
|
||||
# (*, N, D), (*, D), (*, D, D)
|
||||
U, S, Vh = torch.linalg.svd(points - mean)
|
||||
normal = Vh[..., -1, :] # (*, D)
|
||||
offset = torch.einsum("...ij,...j->...i", points, normal) # (*, N)
|
||||
offset = offset.mean(dim=-1, keepdim=True)
|
||||
return torch.cat([normal, offset], dim=-1)
|
||||
|
||||
|
||||
def get_plane_transform(up, ground_plane=None, xyz_orig=None):
|
||||
"""
|
||||
get R, t rigid transform from plane and desired origin
|
||||
:param up (3,) up vector of coordinate frame
|
||||
:param ground_plane (4) (a, b, c, d) where a,b,c is the normal
|
||||
:param xyz_orig (3) desired origin
|
||||
"""
|
||||
R = torch.eye(3)
|
||||
t = torch.zeros(3)
|
||||
if ground_plane is None:
|
||||
return R, t
|
||||
|
||||
# compute transform between world up vector and passed in floor
|
||||
ground_plane = torch.as_tensor(ground_plane)
|
||||
ground_plane = torch.sign(ground_plane[3]) * ground_plane
|
||||
|
||||
normal = ground_plane[:3]
|
||||
normal = normal / torch.linalg.norm(normal)
|
||||
v = torch.linalg.cross(up, normal)
|
||||
ang_sin = torch.linalg.norm(v)
|
||||
ang_cos = up.dot(normal)
|
||||
skew_v = torch.as_tensor([[0.0, -v[2], v[1]], [v[2], 0.0, -v[0]], [-v[1], v[0], 0.0]])
|
||||
R = torch.eye(3) + skew_v + (skew_v @ skew_v) * ((1.0 - ang_cos) / (ang_sin**2))
|
||||
|
||||
# project origin onto plane
|
||||
if xyz_orig is None:
|
||||
xyz_orig = torch.zeros(3)
|
||||
t, _ = compute_plane_intersection(xyz_orig, -normal, ground_plane)
|
||||
|
||||
return R, t
|
||||
|
||||
|
||||
def parse_floor_plane(floor_plane):
|
||||
"""
|
||||
Takes floor plane in the optimization form (Bx3 with a,b,c * d) and parses into
|
||||
(a,b,c,d) from with (a,b,c) normal facing "up in the camera frame and d the offset.
|
||||
"""
|
||||
floor_offset = torch.norm(floor_plane, dim=-1, keepdim=True)
|
||||
floor_normal = floor_plane / floor_offset
|
||||
|
||||
# in camera system -y is up, so floor plane normal y component should never be positive
|
||||
# (assuming the camera is not sideways or upside down)
|
||||
neg_mask = floor_normal[..., 1:2] > 0.0
|
||||
floor_normal = torch.where(
|
||||
neg_mask.expand_as(floor_normal), -floor_normal, floor_normal
|
||||
)
|
||||
floor_offset = torch.where(neg_mask, -floor_offset, floor_offset)
|
||||
floor_plane_4d = torch.cat([floor_normal, floor_offset], dim=-1)
|
||||
|
||||
return floor_plane_4d
|
||||
|
||||
|
||||
def compute_plane_intersection(point, direction, plane):
|
||||
"""
|
||||
Given a ray defined by a point in space and a direction,
|
||||
compute the intersection point with the given plane.
|
||||
Detect intersection in either direction or -direction.
|
||||
Note, ray may not actually intersect with the plane.
|
||||
|
||||
Returns the intersection point and s where
|
||||
point + s * direction = intersection_point. if s < 0 it means
|
||||
-direction intersects.
|
||||
|
||||
- point : B x 3
|
||||
- direction : B x 3
|
||||
- plane : B x 4 (a, b, c, d) where (a, b, c) is the normal and (d) the offset.
|
||||
"""
|
||||
dims = point.shape[:-1]
|
||||
plane_normal = plane[..., :3]
|
||||
plane_off = plane[..., 3]
|
||||
s = (plane_off - bdot(plane_normal, point)) / (bdot(plane_normal, direction) + 1e-4)
|
||||
itsct_pt = point + s.reshape((-1, 1)) * direction
|
||||
return itsct_pt, s
|
||||
|
||||
|
||||
def bdot(A1, A2, keepdim=False):
|
||||
"""
|
||||
Batched dot product.
|
||||
- A1 : B x D
|
||||
- A2 : B x D.
|
||||
Returns B.
|
||||
"""
|
||||
return (A1 * A2).sum(dim=-1, keepdim=keepdim)
|
||||
+293
@@ -0,0 +1,293 @@
|
||||
import torch
|
||||
import numpy as np
|
||||
from torch.nn import functional as F
|
||||
|
||||
|
||||
def batch_rodrigues(rot_vecs, epsilon=1e-8, dtype=torch.float32):
|
||||
"""
|
||||
Taken from https://github.com/mkocabas/VIBE/blob/master/lib/utils/geometry.py
|
||||
Calculates the rotation matrices for a batch of rotation vectors
|
||||
- param rot_vecs: torch.tensor (N, 3) array of N axis-angle vectors
|
||||
- returns R: torch.tensor (N, 3, 3) rotation matrices
|
||||
"""
|
||||
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
|
||||
|
||||
|
||||
def quaternion_mul(q0, q1):
|
||||
"""
|
||||
EXPECTS WXYZ
|
||||
:param q0 (*, 4)
|
||||
:param q1 (*, 4)
|
||||
"""
|
||||
r0, r1 = q0[..., :1], q1[..., :1]
|
||||
v0, v1 = q0[..., 1:], q1[..., 1:]
|
||||
r = r0 * r1 - (v0 * v1).sum(dim=-1, keepdim=True)
|
||||
v = r0 * v1 + r1 * v0 + torch.linalg.cross(v0, v1)
|
||||
return torch.cat([r, v], dim=-1)
|
||||
|
||||
|
||||
def quaternion_inverse(q, eps=1e-8):
|
||||
"""
|
||||
EXPECTS WXYZ
|
||||
:param q (*, 4)
|
||||
"""
|
||||
conj = torch.cat([q[..., :1], -q[..., 1:]], dim=-1)
|
||||
mag = torch.square(q).sum(dim=-1, keepdim=True) + eps
|
||||
return conj / mag
|
||||
|
||||
|
||||
def quaternion_slerp(t, q0, q1, eps=1e-8):
|
||||
"""
|
||||
:param t (*, 1) must be between 0 and 1
|
||||
:param q0 (*, 4)
|
||||
:param q1 (*, 4)
|
||||
"""
|
||||
dims = q0.shape[:-1]
|
||||
t = t.view(*dims, 1)
|
||||
|
||||
q0 = F.normalize(q0, p=2, dim=-1)
|
||||
q1 = F.normalize(q1, p=2, dim=-1)
|
||||
dot = (q0 * q1).sum(dim=-1, keepdim=True)
|
||||
|
||||
# make sure we give the shortest rotation path (< 180d)
|
||||
neg = dot < 0
|
||||
q1 = torch.where(neg, -q1, q1)
|
||||
dot = torch.where(neg, -dot, dot)
|
||||
angle = torch.acos(dot)
|
||||
|
||||
# if angle is too small, just do linear interpolation
|
||||
collin = torch.abs(dot) > 1 - eps
|
||||
fac = 1 / torch.sin(angle)
|
||||
w0 = torch.where(collin, 1 - t, torch.sin((1 - t) * angle) * fac)
|
||||
w1 = torch.where(collin, t, torch.sin(t * angle) * fac)
|
||||
slerp = q0 * w0 + q1 * w1
|
||||
return slerp
|
||||
|
||||
|
||||
def rotation_matrix_to_angle_axis(rotation_matrix):
|
||||
"""
|
||||
This function is borrowed from https://github.com/kornia/kornia
|
||||
|
||||
Convert rotation matrix to Rodrigues vector
|
||||
"""
|
||||
quaternion = rotation_matrix_to_quaternion(rotation_matrix)
|
||||
aa = quaternion_to_angle_axis(quaternion)
|
||||
aa[torch.isnan(aa)] = 0.0
|
||||
return aa
|
||||
|
||||
|
||||
def quaternion_to_angle_axis(quaternion):
|
||||
"""
|
||||
This function is borrowed from https://github.com/kornia/kornia
|
||||
|
||||
Convert quaternion vector to angle axis of rotation.
|
||||
Adapted from ceres C++ library: ceres-solver/include/ceres/rotation.h
|
||||
|
||||
:param quaternion (*, 4) expects WXYZ
|
||||
:returns angle_axis (*, 3)
|
||||
"""
|
||||
# unpack input and compute conversion
|
||||
q1 = quaternion[..., 1]
|
||||
q2 = quaternion[..., 2]
|
||||
q3 = quaternion[..., 3]
|
||||
sin_squared_theta = q1 * q1 + q2 * q2 + q3 * q3
|
||||
|
||||
sin_theta = torch.sqrt(sin_squared_theta)
|
||||
cos_theta = quaternion[..., 0]
|
||||
two_theta = 2.0 * torch.where(
|
||||
cos_theta < 0.0,
|
||||
torch.atan2(-sin_theta, -cos_theta),
|
||||
torch.atan2(sin_theta, cos_theta),
|
||||
)
|
||||
|
||||
k_pos = two_theta / sin_theta
|
||||
k_neg = 2.0 * torch.ones_like(sin_theta)
|
||||
k = torch.where(sin_squared_theta > 0.0, k_pos, k_neg)
|
||||
|
||||
angle_axis = torch.zeros_like(quaternion)[..., :3]
|
||||
angle_axis[..., 0] += q1 * k
|
||||
angle_axis[..., 1] += q2 * k
|
||||
angle_axis[..., 2] += q3 * k
|
||||
return angle_axis
|
||||
|
||||
|
||||
def angle_axis_to_rotation_matrix(angle_axis):
|
||||
"""
|
||||
:param angle_axis (*, 3)
|
||||
return (*, 3, 3)
|
||||
"""
|
||||
quat = angle_axis_to_quaternion(angle_axis)
|
||||
return quaternion_to_rotation_matrix(quat)
|
||||
|
||||
|
||||
def quaternion_to_rotation_matrix(quaternion):
|
||||
"""
|
||||
Convert a quaternion to a rotation matrix.
|
||||
Taken from https://github.com/kornia/kornia, based on
|
||||
https://github.com/matthew-brett/transforms3d/blob/8965c48401d9e8e66b6a8c37c65f2fc200a076fa/transforms3d/quaternions.py#L101
|
||||
https://github.com/tensorflow/graphics/blob/master/tensorflow_graphics/geometry/transformation/rotation_matrix_3d.py#L247
|
||||
:param quaternion (N, 4) expects WXYZ order
|
||||
returns rotation matrix (N, 3, 3)
|
||||
"""
|
||||
# normalize the input quaternion
|
||||
quaternion_norm = F.normalize(quaternion, p=2, dim=-1, eps=1e-12)
|
||||
*dims, _ = quaternion_norm.shape
|
||||
|
||||
# unpack the normalized quaternion components
|
||||
w, x, y, z = torch.chunk(quaternion_norm, chunks=4, dim=-1)
|
||||
|
||||
# compute the actual conversion
|
||||
tx = 2.0 * x
|
||||
ty = 2.0 * y
|
||||
tz = 2.0 * z
|
||||
twx = tx * w
|
||||
twy = ty * w
|
||||
twz = tz * w
|
||||
txx = tx * x
|
||||
txy = ty * x
|
||||
txz = tz * x
|
||||
tyy = ty * y
|
||||
tyz = tz * y
|
||||
tzz = tz * z
|
||||
one = torch.tensor(1.0)
|
||||
|
||||
matrix = torch.stack(
|
||||
(
|
||||
one - (tyy + tzz),
|
||||
txy - twz,
|
||||
txz + twy,
|
||||
txy + twz,
|
||||
one - (txx + tzz),
|
||||
tyz - twx,
|
||||
txz - twy,
|
||||
tyz + twx,
|
||||
one - (txx + tyy),
|
||||
),
|
||||
dim=-1,
|
||||
).view(*dims, 3, 3)
|
||||
return matrix
|
||||
|
||||
|
||||
def angle_axis_to_quaternion(angle_axis):
|
||||
"""
|
||||
This function is borrowed from https://github.com/kornia/kornia
|
||||
Convert angle axis to quaternion in WXYZ order
|
||||
:param angle_axis (*, 3)
|
||||
:returns quaternion (*, 4) WXYZ order
|
||||
"""
|
||||
theta_sq = torch.sum(angle_axis**2, dim=-1, keepdim=True) # (*, 1)
|
||||
# need to handle the zero rotation case
|
||||
valid = theta_sq > 0
|
||||
theta = torch.sqrt(theta_sq)
|
||||
half_theta = 0.5 * theta
|
||||
ones = torch.ones_like(half_theta)
|
||||
# fill zero with the limit of sin ax / x -> a
|
||||
k = torch.where(valid, torch.sin(half_theta) / theta, 0.5 * ones)
|
||||
w = torch.where(valid, torch.cos(half_theta), ones)
|
||||
quat = torch.cat([w, k * angle_axis], dim=-1)
|
||||
return quat
|
||||
|
||||
|
||||
def rotation_matrix_to_quaternion(rotation_matrix, eps=1e-6):
|
||||
"""
|
||||
This function is borrowed from https://github.com/kornia/kornia
|
||||
Convert rotation matrix to 4d quaternion vector
|
||||
This algorithm is based on algorithm described in
|
||||
https://github.com/KieranWynn/pyquaternion/blob/master/pyquaternion/quaternion.py#L201
|
||||
|
||||
:param rotation_matrix (N, 3, 3)
|
||||
"""
|
||||
*dims, m, n = rotation_matrix.shape
|
||||
rmat_t = torch.transpose(rotation_matrix.reshape(-1, m, n), -1, -2)
|
||||
|
||||
mask_d2 = rmat_t[:, 2, 2] < eps
|
||||
|
||||
mask_d0_d1 = rmat_t[:, 0, 0] > rmat_t[:, 1, 1]
|
||||
mask_d0_nd1 = rmat_t[:, 0, 0] < -rmat_t[:, 1, 1]
|
||||
|
||||
t0 = 1 + rmat_t[:, 0, 0] - rmat_t[:, 1, 1] - rmat_t[:, 2, 2]
|
||||
q0 = torch.stack(
|
||||
[
|
||||
rmat_t[:, 1, 2] - rmat_t[:, 2, 1],
|
||||
t0,
|
||||
rmat_t[:, 0, 1] + rmat_t[:, 1, 0],
|
||||
rmat_t[:, 2, 0] + rmat_t[:, 0, 2],
|
||||
],
|
||||
-1,
|
||||
)
|
||||
t0_rep = t0.repeat(4, 1).t()
|
||||
|
||||
t1 = 1 - rmat_t[:, 0, 0] + rmat_t[:, 1, 1] - rmat_t[:, 2, 2]
|
||||
q1 = torch.stack(
|
||||
[
|
||||
rmat_t[:, 2, 0] - rmat_t[:, 0, 2],
|
||||
rmat_t[:, 0, 1] + rmat_t[:, 1, 0],
|
||||
t1,
|
||||
rmat_t[:, 1, 2] + rmat_t[:, 2, 1],
|
||||
],
|
||||
-1,
|
||||
)
|
||||
t1_rep = t1.repeat(4, 1).t()
|
||||
|
||||
t2 = 1 - rmat_t[:, 0, 0] - rmat_t[:, 1, 1] + rmat_t[:, 2, 2]
|
||||
q2 = torch.stack(
|
||||
[
|
||||
rmat_t[:, 0, 1] - rmat_t[:, 1, 0],
|
||||
rmat_t[:, 2, 0] + rmat_t[:, 0, 2],
|
||||
rmat_t[:, 1, 2] + rmat_t[:, 2, 1],
|
||||
t2,
|
||||
],
|
||||
-1,
|
||||
)
|
||||
t2_rep = t2.repeat(4, 1).t()
|
||||
|
||||
t3 = 1 + rmat_t[:, 0, 0] + rmat_t[:, 1, 1] + rmat_t[:, 2, 2]
|
||||
q3 = torch.stack(
|
||||
[
|
||||
t3,
|
||||
rmat_t[:, 1, 2] - rmat_t[:, 2, 1],
|
||||
rmat_t[:, 2, 0] - rmat_t[:, 0, 2],
|
||||
rmat_t[:, 0, 1] - rmat_t[:, 1, 0],
|
||||
],
|
||||
-1,
|
||||
)
|
||||
t3_rep = t3.repeat(4, 1).t()
|
||||
|
||||
mask_c0 = mask_d2 * mask_d0_d1
|
||||
mask_c1 = mask_d2 * ~mask_d0_d1
|
||||
mask_c2 = ~mask_d2 * mask_d0_nd1
|
||||
mask_c3 = ~mask_d2 * ~mask_d0_nd1
|
||||
mask_c0 = mask_c0.view(-1, 1).type_as(q0)
|
||||
mask_c1 = mask_c1.view(-1, 1).type_as(q1)
|
||||
mask_c2 = mask_c2.view(-1, 1).type_as(q2)
|
||||
mask_c3 = mask_c3.view(-1, 1).type_as(q3)
|
||||
|
||||
q = q0 * mask_c0 + q1 * mask_c1 + q2 * mask_c2 + q3 * mask_c3
|
||||
q /= torch.sqrt(
|
||||
t0_rep * mask_c0
|
||||
+ t1_rep * mask_c1
|
||||
+ t2_rep * mask_c2 # noqa
|
||||
+ t3_rep * mask_c3
|
||||
) # noqa
|
||||
q *= 0.5
|
||||
return q.reshape(*dims, 4)
|
||||
+309
@@ -0,0 +1,309 @@
|
||||
import os
|
||||
import numpy as np
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from body_model import MANO_JOINTS, run_mano
|
||||
from geometry.rotation import (
|
||||
rotation_matrix_to_angle_axis,
|
||||
angle_axis_to_rotation_matrix,
|
||||
)
|
||||
from util.logger import Logger
|
||||
from util.tensor import move_to, detach_all
|
||||
|
||||
# from .helpers import estimate_initial_trans
|
||||
from .params import CameraParams
|
||||
|
||||
J_HAND = len(MANO_JOINTS) - 1 # no root
|
||||
|
||||
|
||||
class BaseSceneModel(nn.Module):
|
||||
"""
|
||||
Scene model of sequences of human poses.
|
||||
All poses are in their own INDEPENDENT camera reference frames.
|
||||
A basic class mostly for testing purposes.
|
||||
|
||||
Parameters:
|
||||
batch_size: number of sequences to optimize
|
||||
seq_len: length of the sequences
|
||||
body_model: MANO hand model
|
||||
pose_prior: VPoser model
|
||||
fit_gender: gender of model (optional)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
batch_size,
|
||||
seq_len,
|
||||
body_model,
|
||||
pose_prior,
|
||||
# fit_gender="neutral",
|
||||
use_init=False,
|
||||
opt_cams=False,
|
||||
opt_scale=True,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__()
|
||||
B, T = batch_size, seq_len
|
||||
self.batch_size = batch_size
|
||||
self.seq_len = seq_len
|
||||
|
||||
self.body_model = body_model
|
||||
self.hand_mean = self.body_model.hand_mean
|
||||
|
||||
self.pose_prior = pose_prior
|
||||
if self.pose_prior is not None:
|
||||
self.latent_pose_dim = self.pose_prior.latentD
|
||||
|
||||
self.num_betas = body_model.num_betas
|
||||
|
||||
self.use_init = use_init
|
||||
print("USE INIT", use_init)
|
||||
self.opt_scale = opt_scale
|
||||
self.opt_cams = opt_cams
|
||||
print("OPT SCALE", self.opt_scale)
|
||||
print("OPT CAMERAS", self.opt_cams)
|
||||
print("Batch size: ", batch_size)
|
||||
self.params = CameraParams(batch_size)
|
||||
|
||||
def initialize(self, obs_data, cam_data):
|
||||
Logger.log("Initializing scene model with observed data")
|
||||
|
||||
# initialize cameras
|
||||
self.params.set_cameras(
|
||||
cam_data,
|
||||
opt_scale=self.opt_scale,
|
||||
opt_cams=self.opt_cams,
|
||||
opt_focal=self.opt_cams,
|
||||
)
|
||||
|
||||
# initialize body params
|
||||
B, T = self.batch_size, self.seq_len
|
||||
device = next(iter(cam_data.values())).device
|
||||
init_betas = torch.mean(obs_data["init_body_shape"], dim=1) # torch.zeros(B, self.num_betas, device=device)
|
||||
|
||||
if self.use_init and "init_body_pose" in obs_data:
|
||||
init_pose = obs_data["init_body_pose"][:, :, :J_HAND, :]
|
||||
init_pose_latent = self.pose2latent(init_pose)
|
||||
else:
|
||||
raise ValueError
|
||||
# init_pose = torch.zeros(B, T, J_HAND, 3, device=device)
|
||||
# init_pose_latent = torch.zeros(B, T, self.latent_pose_dim, device=device)
|
||||
|
||||
# transform into world frame (T, 3, 3), (T, 3)
|
||||
R_w2c, t_w2c = cam_data["cam_R"], cam_data["cam_t"]
|
||||
R_c2w = R_w2c.transpose(-1, -2)
|
||||
t_c2w = -torch.einsum("tij,tj->ti", R_c2w, t_w2c)
|
||||
|
||||
if self.use_init and "init_root_orient" in obs_data:
|
||||
init_rot = obs_data["init_root_orient"] # (B, T, 3)
|
||||
init_rot_mat = angle_axis_to_rotation_matrix(init_rot)
|
||||
print(R_c2w.shape, init_rot_mat.shape)
|
||||
init_rot_mat = torch.einsum("tij,btjk->btik", R_c2w, init_rot_mat)
|
||||
init_rot = rotation_matrix_to_angle_axis(init_rot_mat)
|
||||
else:
|
||||
raise ValueError
|
||||
# init_rot = (
|
||||
# torch.tensor([np.pi, 0, 0], dtype=torch.float32)
|
||||
# .reshape(1, 1, 3)
|
||||
# .repeat(B, T, 1)
|
||||
# )
|
||||
|
||||
init_trans = torch.zeros(B, T, 3, device=device)
|
||||
if self.use_init and "init_trans" in obs_data:
|
||||
# must offset by the root location before applying camera to world transform
|
||||
is_right = obs_data["is_right"]
|
||||
pred_data = self.pred_mano(init_trans, init_rot, init_pose, is_right, init_betas)
|
||||
root_loc = pred_data["joints3d"][..., 0, :] # (B, T, 3)
|
||||
init_trans = obs_data["init_trans"] # (B, T, 3)
|
||||
init_trans = (
|
||||
torch.einsum("tij,btj->bti", R_c2w, init_trans + root_loc)
|
||||
+ t_c2w[None]
|
||||
- root_loc
|
||||
)
|
||||
else:
|
||||
# initialize trans with reprojected joints
|
||||
# pred_data = self.pred_mano(init_trans, init_rot, init_pose, init_betas)
|
||||
# init_trans = estimate_initial_trans(
|
||||
# init_pose,
|
||||
# pred_data["joints3d_op"],
|
||||
# obs_data["joints2d"],
|
||||
# obs_data["intrins"][:, 0],
|
||||
# )
|
||||
raise ValueError
|
||||
|
||||
self.params.set_param("init_body_pose", init_pose)
|
||||
self.params.set_param("latent_pose", init_pose_latent)
|
||||
self.params.set_param("betas", init_betas)
|
||||
self.params.set_param("trans", init_trans)
|
||||
self.params.set_param("root_orient", init_rot)
|
||||
self.params.set_param("is_right", is_right, requires_grad=False)
|
||||
|
||||
# Store initial latent pose in obs_data for pose prior loss
|
||||
obs_data["init_latent_pose"] = init_pose_latent.detach()
|
||||
# print(init_pose.shape, init_pose_latent.shape, init_betas.shape, init_trans.shape, init_rot.shape, is_right.shape)
|
||||
# raise ValueError
|
||||
|
||||
def get_optim_result(self, **kwargs):
|
||||
"""
|
||||
Collect predicted outputs (latent_pose, trans, root_orient, betas, body pose) into dict
|
||||
"""
|
||||
print('running get_optim_result in base_scene.py...')
|
||||
res = self.params.get_dict()
|
||||
if "latent_pose" in res:
|
||||
res["pose_body"] = self.latent2pose(self.params.latent_pose).detach()
|
||||
|
||||
# add the cameras
|
||||
res["cam_R"], res["cam_t"], _, _ = self.params.get_cameras()
|
||||
res["intrins"] = self.params.intrins
|
||||
return {"world": res}
|
||||
|
||||
def latent2pose(self, latent_pose):
|
||||
"""
|
||||
Converts VPoser latent embedding to aa body pose.
|
||||
latent_pose : B x T x D
|
||||
body_pose : B x T x J*3
|
||||
"""
|
||||
if self.pose_prior is not None:
|
||||
B, T, _ = latent_pose.size()
|
||||
d_latent = self.pose_prior.latentD
|
||||
latent_pose = latent_pose.reshape((-1, d_latent))
|
||||
body_pose = self.pose_prior.decode(latent_pose, output_type="matrot")
|
||||
body_pose = rotation_matrix_to_angle_axis(
|
||||
body_pose.reshape((B * T * J_HAND, 3, 3))
|
||||
).reshape((B, T, J_HAND * 3))
|
||||
return body_pose + self.hand_mean
|
||||
else:
|
||||
return latent_pose
|
||||
|
||||
def pose2latent(self, body_pose):
|
||||
"""
|
||||
Encodes aa body pose to VPoser latent space.
|
||||
body_pose : B x T x J*3
|
||||
latent_pose : B x T x D
|
||||
"""
|
||||
if self.pose_prior is not None:
|
||||
B, T = body_pose.shape[:2]
|
||||
body_pose = body_pose.reshape((-1, J_HAND * 3))
|
||||
latent_pose_distrib = self.pose_prior.encode(body_pose - self.hand_mean)
|
||||
d_latent = self.pose_prior.latentD
|
||||
latent_pose = latent_pose_distrib.mean.reshape((B, T, d_latent))
|
||||
return latent_pose
|
||||
else:
|
||||
return body_pose
|
||||
|
||||
def pred_mano(self, trans, root_orient, body_pose, is_right, betas):
|
||||
"""
|
||||
Forward pass of the MANO model and populates pred_data accordingly with
|
||||
joints3d, verts3d, points3d.
|
||||
|
||||
trans : B x T x 3
|
||||
root_orient : B x T x 3
|
||||
body_pose : B x T x J*3
|
||||
betas : B x D
|
||||
"""
|
||||
|
||||
mano_out = run_mano(self.body_model, trans, root_orient, body_pose, is_right, betas=betas)
|
||||
joints3d, points3d = mano_out["joints"], mano_out["vertices"]
|
||||
|
||||
# select desired joints and vertices
|
||||
# joints3d_body = joints3d
|
||||
# joints3d_op = joints3d.clone()
|
||||
# joints3d_op = joints3d[:, :, self.smpl2op_map, :]
|
||||
# # hacky way to get hip joints that align with ViTPose keypoints
|
||||
# # this could be moved elsewhere in the future (and done properly)
|
||||
# joints3d_op[:, :, [9, 12]] = (
|
||||
# joints3d_op[:, :, [9, 12]]
|
||||
# + 0.25 * (joints3d_op[:, :, [9, 12]] - joints3d_op[:, :, [12, 9]])
|
||||
# + 0.5
|
||||
# * (
|
||||
# joints3d_op[:, :, [8]]
|
||||
# - 0.5 * (joints3d_op[:, :, [9, 12]] + joints3d_op[:, :, [12, 9]])
|
||||
# )
|
||||
# )
|
||||
verts3d = points3d #[:, :, KEYPT_VERTS, :]
|
||||
|
||||
return {
|
||||
"is_right": is_right,
|
||||
"points3d": points3d, # all vertices
|
||||
"verts3d": verts3d, # keypoint vertices
|
||||
"joints3d": joints3d, # smpl joints
|
||||
"joints3d_op": joints3d, # OP joints
|
||||
"l_faces": mano_out["l_faces"], # index array of faces
|
||||
"r_faces": mano_out["r_faces"], # index array of faces
|
||||
"body_pose": mano_out["body_pose"]
|
||||
}
|
||||
|
||||
def pred_params_mano(self, is_right, reproj=True):
|
||||
body_pose = self.latent2pose(self.params.latent_pose)
|
||||
pred_data = self.pred_mano(
|
||||
self.params.trans, self.params.root_orient, body_pose, is_right, self.params.betas
|
||||
)
|
||||
# from geometry.mesh import save_mesh_scenes, vertices_to_trimesh
|
||||
# LIGHT_BLUE=(0.65098039, 0.74117647, 0.85882353)
|
||||
# print(pred_data.keys())
|
||||
# verts, l_faces, r_faces, is_right = pred_data['verts3d'], pred_data['l_faces'], pred_data['l_faces'], pred_data['is_right']
|
||||
# T = len(verts[0])
|
||||
# print(f"{T} mesh frames")
|
||||
# print(verts.shape, l_faces.shape, is_right.shape, is_right)
|
||||
# times = list(range(0, T, 1))
|
||||
# os.makedirs('./mesh/')
|
||||
# for t in times:
|
||||
|
||||
# tmesh = vertices_to_trimesh(verts[0][t].detach().cpu().numpy(), l_faces.detach().cpu().numpy(), LIGHT_BLUE, is_right=0)
|
||||
# tmesh.export(os.path.join('./mesh/', f'{str(t).zfill(6)}_0.obj'))
|
||||
|
||||
# tmesh = vertices_to_trimesh(verts[1][t].detach().cpu().numpy(), r_faces.detach().cpu().numpy(), LIGHT_BLUE, is_right=1)
|
||||
# tmesh.export(os.path.join('./mesh/', f'{str(t).zfill(6)}_1.obj'))
|
||||
# print(os.path.join('./mesh/', f'{str(t).zfill(6)}_1.obj'))
|
||||
|
||||
#################################
|
||||
# print(self.params.root_orient.shape, self.params.init_body_pose.shape, body_pose.shape, self.params.trans.shape, self.params.betas.shape)
|
||||
# hand_model = mano.load(
|
||||
# model_path= '/vol/bitbucket/zy3023/code/hand/slam-hand/_DATA/data/mano',# '/vol/bitbucket/zy3023/code/hand/slam-hand/VPoser/pretrained/Vposer_right_mirrored',
|
||||
# is_right=True,
|
||||
# num_pca_comps=45,
|
||||
# batch_size=1,
|
||||
# flat_hand_mean=False,
|
||||
# ).cuda()
|
||||
|
||||
# output = hand_model(
|
||||
# betas=self.params.betas[0][None].repeat(180,1),
|
||||
# global_orient=self.params.root_orient[0],
|
||||
# hand_pose=self.params.init_body_pose[0].reshape(180, -1) - self.hand_mean,
|
||||
# transl=self.params.trans[0],
|
||||
# return_verts=True,
|
||||
# return_tips=False
|
||||
# )
|
||||
# np.save('x.npy', self.params.init_body_pose[0].reshape(180, -1).detach().cpu().numpy())
|
||||
|
||||
# out_path = './out'
|
||||
# h_meshes = hand_model.hand_meshes(output)
|
||||
# if not os.path.exists(out_path):
|
||||
# os.mkdir(out_path + '_after')
|
||||
# os.mkdir(out_path)
|
||||
|
||||
# for i in range(180):
|
||||
# #print(f"export to {out_path}/{str(i).zfill(6)}.ply")
|
||||
# h_meshes[i].export(f"{out_path}/{str(i).zfill(6)}.ply")
|
||||
|
||||
# # print(body_pose.shape, body_pose[0])
|
||||
# print(self.hand_mean)
|
||||
# output = hand_model(
|
||||
# betas=self.params.betas[0][None].repeat(180,1),
|
||||
# global_orient=self.params.root_orient[0],
|
||||
# hand_pose=body_pose[0].reshape(180, -1) - self.hand_mean,
|
||||
# transl=self.params.trans[0],
|
||||
# return_verts=True,
|
||||
# return_tips=False
|
||||
# )
|
||||
|
||||
# h_meshes = hand_model.hand_meshes(output)
|
||||
# for i in range(180):
|
||||
# #print(f"export to {out_path}_after/{str(i).zfill(6)}.ply")
|
||||
# h_meshes[i].export(f"{out_path}_after/{str(i).zfill(6)}.ply")
|
||||
# exit()
|
||||
|
||||
return pred_data
|
||||
+399
@@ -0,0 +1,399 @@
|
||||
import os
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as torch_f
|
||||
|
||||
SNAP_PARENT = [
|
||||
0, # 0's parent
|
||||
0, # 1's parent
|
||||
1,
|
||||
2,
|
||||
3,
|
||||
0, # 5's parent
|
||||
5,
|
||||
6,
|
||||
7,
|
||||
0, # 9's parent
|
||||
9,
|
||||
10,
|
||||
11,
|
||||
0, # 13's parent
|
||||
13,
|
||||
14,
|
||||
15,
|
||||
0, # 17's parent
|
||||
17,
|
||||
18,
|
||||
19,
|
||||
]
|
||||
|
||||
JOINT_ROOT_IDX = 9
|
||||
|
||||
REF_BONE_LINK = (0, 9) # mid mcp
|
||||
|
||||
# bone indexes in 20 bones setting
|
||||
ID_ROOT_bone = np.array([0, 4, 8, 12, 16]) # ROOT_bone from wrist to MCP
|
||||
ID_PIP_bone = np.array([1, 5, 9, 13, 17]) # PIP_bone from MCP to PIP
|
||||
ID_DIP_bone = np.array([2, 6, 10, 14, 18]) # DIP_bone from PIP to DIP
|
||||
ID_TIP_bone = np.array([3, 7, 11, 15, 19]) # TIP_bone from DIP to TIP
|
||||
|
||||
def plot_hull(theta, hull):
|
||||
del_rdp_hull = hull.detach().cpu().numpy()
|
||||
theta = theta.detach().cpu().numpy()
|
||||
|
||||
fig = plt.figure()
|
||||
figManager = plt.get_current_fig_manager()
|
||||
figManager.window.showMaximized()
|
||||
|
||||
ax = fig.add_subplot(111)
|
||||
ax.scatter(theta[:, 0], theta[:, 1], s=10, c='r')
|
||||
ax.set_xlabel("flexion")
|
||||
ax.set_ylabel("abduction")
|
||||
|
||||
plt.plot(del_rdp_hull[:, 0],
|
||||
del_rdp_hull[:, 1],
|
||||
'-yo', linewidth=2)
|
||||
|
||||
plt.xticks(np.arange(-3, 4, 1))
|
||||
plt.yticks(np.arange(-2, 2, 0.5))
|
||||
plt.show()
|
||||
|
||||
|
||||
def two_norm(a):
|
||||
'''
|
||||
|
||||
Args:
|
||||
a: B*M*2 or B*M*3
|
||||
|
||||
Returns:
|
||||
|
||||
'''
|
||||
return torch.norm(a, dim=-1)
|
||||
|
||||
|
||||
def one_norm(a):
|
||||
'''
|
||||
|
||||
Args:
|
||||
a: B*M*2 or B*M*3
|
||||
|
||||
Returns:
|
||||
|
||||
'''
|
||||
return torch.norm(a, dim=-1, p=1)
|
||||
|
||||
|
||||
def calculate_joint_angle_loss(thetas, hulls):
|
||||
'''
|
||||
|
||||
Args:
|
||||
Theta: B*15*2
|
||||
hulls: list
|
||||
|
||||
Returns:
|
||||
|
||||
'''
|
||||
|
||||
loss = torch.Tensor([0]).cuda()
|
||||
for i in range(15):
|
||||
# print("i=",i)
|
||||
hull = hulls[i] # (M*2)
|
||||
theta = thetas[:, i] # (B*2)
|
||||
hull = torch.cat((hull, hull[0].unsqueeze(0)), dim=0)
|
||||
|
||||
v = (hull[1:] - hull[:-1]).unsqueeze(0) # (M-1)*2
|
||||
w = - hull[:-1].unsqueeze(0) + theta.unsqueeze(1).repeat(1, hull[:-1].shape[0], 1) # B*(M-1)*2
|
||||
|
||||
cross_product_2d = w[:, :, 0] * v[:, :, 1] - w[:, :, 1] * v[:, :, 0]
|
||||
tmp = torch.sum(cross_product_2d < 1e-6, dim=-1)
|
||||
|
||||
is_outside = (tmp != (hull.shape[0] - 1))
|
||||
if not torch.sum(is_outside):
|
||||
sub_loss = torch.Tensor([0]).cuda()
|
||||
else:
|
||||
outside_theta = theta[is_outside]
|
||||
outside_theta = outside_theta.unsqueeze(1).repeat(1, hull[:-1].shape[0], 1)
|
||||
w_outside = - hull[:-1].unsqueeze(0) + outside_theta # B*(M-1)*2
|
||||
t = torch.clamp(inner_product(w_outside, v) / (two_norm(v) ** 2), min=0, max=1).unsqueeze(2)
|
||||
p = hull[:-1] + t * v
|
||||
|
||||
D = one_norm(torch.cos(outside_theta) - torch.cos(p)) + one_norm(torch.sin(outside_theta) - torch.sin(p))
|
||||
sub_loss = torch.sum(torch.min(D, dim=-1)[0])
|
||||
|
||||
vis = 0
|
||||
if vis:
|
||||
print(theta)
|
||||
plot_hull(theta, hull)
|
||||
|
||||
loss += sub_loss
|
||||
|
||||
loss /= (15 * thetas.shape[0])
|
||||
|
||||
return loss
|
||||
|
||||
|
||||
def angle_between(v1, v2):
|
||||
epsilon = 1e-7
|
||||
cos = torch_f.cosine_similarity(v1, v2, dim=-1).clamp(-1 + epsilon, 1 - epsilon) # (B)
|
||||
theta = torch.acos(cos) # (B)
|
||||
return theta
|
||||
|
||||
|
||||
def normalize(vec):
|
||||
return torch_f.normalize(vec, p=2, dim=-1)
|
||||
|
||||
|
||||
def inner_product(x1, x2):
|
||||
return torch.sum(x1 * x2, dim=-1)
|
||||
|
||||
|
||||
def cross_product(x1, x2):
|
||||
return torch.cross(x1, x2, dim=-1)
|
||||
|
||||
|
||||
def axangle2mat_torch(axis, angle, is_normalized=False):
|
||||
""" Rotation matrix for rotation angle `angle` around `axis`
|
||||
Parameters
|
||||
----------
|
||||
axis : [B,M, 3] element sequence
|
||||
vector specifying axis for rotation.
|
||||
angle :[B,M, ] scalar
|
||||
angle of rotation in radians.
|
||||
is_normalized : bool, optional
|
||||
True if `axis` is already normalized (has norm of 1). Default False.
|
||||
Returns
|
||||
-------
|
||||
mat : array shape (B, M,3,3)
|
||||
rotation matrix for specified rotation
|
||||
Notes
|
||||
-----
|
||||
From: http://en.wikipedia.org/wiki/Rotation_matrix#Axis_and_angle
|
||||
"""
|
||||
B = axis.shape[0]
|
||||
M = axis.shape[1]
|
||||
|
||||
if not is_normalized:
|
||||
norm_axis = axis.norm(p=2, dim=-1, keepdim=True)
|
||||
normed_axis = axis / norm_axis
|
||||
else:
|
||||
normed_axis = axis
|
||||
x, y, z = normed_axis[:, :, 0], normed_axis[:, :, 1], normed_axis[:, :, 2]
|
||||
c = torch.cos(angle)
|
||||
s = torch.sin(angle)
|
||||
C = 1 - c
|
||||
|
||||
xs = x * s;
|
||||
ys = y * s;
|
||||
zs = z * s # noqa
|
||||
xC = x * C;
|
||||
yC = y * C;
|
||||
zC = z * C # noqa
|
||||
xyC = x * yC;
|
||||
yzC = y * zC;
|
||||
zxC = z * xC # noqa
|
||||
|
||||
TMP = torch.stack([x * xC + c, xyC - zs, zxC + ys, xyC + zs, y * yC + c, yzC - xs, zxC - ys, yzC + xs, z * zC + c],
|
||||
dim=-1)
|
||||
return TMP.reshape(B, M, 3, 3)
|
||||
|
||||
|
||||
def interval_loss(value, min, max):
|
||||
'''
|
||||
calculate interval loss
|
||||
Args:
|
||||
value: B*M
|
||||
max: M
|
||||
min: M
|
||||
|
||||
Returns:
|
||||
|
||||
'''
|
||||
|
||||
batch_3d_size = value.shape[0]
|
||||
|
||||
min = min.repeat(value.shape[0], 1)
|
||||
max = max.repeat(value.shape[0], 1)
|
||||
|
||||
loss1 = torch.max(min - value, torch.Tensor([0]).cuda())
|
||||
loss2 = torch.max(value - max, torch.Tensor([0]).cuda())
|
||||
|
||||
loss = (loss1 + loss2).sum()
|
||||
|
||||
loss /= (batch_3d_size * value.shape[1])
|
||||
|
||||
return loss
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
class BMCLoss:
|
||||
def __init__(
|
||||
self,
|
||||
lambda_bl=0.,
|
||||
lambda_rb=0.,
|
||||
lambda_a=0.,
|
||||
):
|
||||
self.lambda_bl = lambda_bl
|
||||
self.lambda_rb = lambda_rb
|
||||
self.lambda_a = lambda_a
|
||||
|
||||
self.lp = os.path.join(os.path.dirname(__file__), "../../_DATA/BMC")
|
||||
# resolved_path = os.path.abspath(self.lp)
|
||||
# print(f"Resolved path: {resolved_path}")
|
||||
# print(f"Current working directory: {os.getcwd()}")
|
||||
# exit()
|
||||
|
||||
self.bone_len_max = np.load(os.path.join(self.lp, "bone_len_max.npy"))
|
||||
self.bone_len_min = np.load(os.path.join(self.lp, "bone_len_min.npy"))
|
||||
self.rb_curvatures_max = np.load(os.path.join(self.lp, "curvatures_max.npy"))
|
||||
self.rb_curvatures_min = np.load(os.path.join(self.lp, "curvatures_min.npy"))
|
||||
self.rb_PHI_max = np.load(os.path.join(self.lp, "PHI_max.npy"))
|
||||
self.rb_PHI_min = np.load(os.path.join(self.lp, "PHI_min.npy"))
|
||||
|
||||
self.joint_angle_limit = np.load(os.path.join(self.lp, "CONVEX_HULLS.npy"),
|
||||
allow_pickle=True)
|
||||
LEN_joint_angle_limit = len(self.joint_angle_limit)
|
||||
|
||||
self.bl_max = torch.from_numpy(self.bone_len_max).float().cuda()
|
||||
self.bl_min = torch.from_numpy(self.bone_len_min).float().cuda()
|
||||
self.rb_curvatures_max = torch.from_numpy(self.rb_curvatures_max).float().cuda()
|
||||
self.rb_curvatures_min = torch.from_numpy(self.rb_curvatures_min).float().cuda()
|
||||
self.rb_PHI_max = torch.from_numpy(self.rb_PHI_max).float().cuda()
|
||||
self.rb_PHI_min = torch.from_numpy(self.rb_PHI_min).float().cuda()
|
||||
|
||||
self.joint_angle_limit = [torch.from_numpy(self.joint_angle_limit[i]).float().cuda() for i in
|
||||
range(LEN_joint_angle_limit)]
|
||||
|
||||
def compute_loss(self, ori_joints, valid_mask):
|
||||
'''
|
||||
|
||||
Args:
|
||||
joints: B*21*3
|
||||
|
||||
Returns:
|
||||
|
||||
'''
|
||||
|
||||
joints = ori_joints[valid_mask]
|
||||
print(joints.shape, ori_joints.shape)
|
||||
batch_size = joints.shape[0]
|
||||
final_loss = torch.Tensor([0]).cuda()
|
||||
|
||||
BMC_losses = {"bmc_bl": torch.Tensor([0]).cuda(), "bmc_rb": torch.Tensor([0]).cuda(), "bmc_a": torch.Tensor([0]).cuda()}
|
||||
|
||||
if (self.lambda_bl < 1e-6) and (self.lambda_rb < 1e-6) and (self.lambda_a < 1e-6):
|
||||
return final_loss, BMC_losses
|
||||
|
||||
ALL_bones = [
|
||||
(
|
||||
joints[:, i, :] -
|
||||
joints[:, SNAP_PARENT[i], :]
|
||||
) for i in range(21)
|
||||
]
|
||||
ALL_bones = torch.stack(ALL_bones[1:], dim=1) # (B,20,3)
|
||||
ROOT_bones = ALL_bones[:, ID_ROOT_bone] # (B,5,3)
|
||||
PIP_bones = ALL_bones[:, ID_PIP_bone]
|
||||
DIP_bones = ALL_bones[:, ID_DIP_bone]
|
||||
TIP_bones = ALL_bones[:, ID_TIP_bone]
|
||||
|
||||
ALL_Z_axis = normalize(ALL_bones)
|
||||
PIP_Z_axis = ALL_Z_axis[:, ID_ROOT_bone]
|
||||
DIP_Z_axis = ALL_Z_axis[:, ID_PIP_bone]
|
||||
TIP_Z_axis = ALL_Z_axis[:, ID_DIP_bone]
|
||||
|
||||
normals = normalize(cross_product(ROOT_bones[:, 1:], ROOT_bones[:, :-1]))
|
||||
|
||||
# compute loss of bone length
|
||||
bl_loss = torch.Tensor([0]).cuda()
|
||||
if self.lambda_bl:
|
||||
bls = two_norm(ALL_bones) # (B,20,1)
|
||||
bl_loss = interval_loss(value=bls, min=self.bl_min, max=self.bl_max)
|
||||
final_loss += self.lambda_bl * bl_loss
|
||||
BMC_losses["bmc_bl"] = bl_loss
|
||||
|
||||
# compute loss of Root bones
|
||||
rb_loss = torch.Tensor([0]).cuda()
|
||||
if self.lambda_rb:
|
||||
edge_normals = torch.zeros_like(ROOT_bones).cuda() # (B,5,3)
|
||||
edge_normals[:, [0, 4]] = normals[:, [0, 3]]
|
||||
edge_normals[:, 1:4] = normalize(normals[:, 1:4] + normals[:, :3])
|
||||
|
||||
curvatures = inner_product(edge_normals[:, 1:] - edge_normals[:, :4],
|
||||
ROOT_bones[:, 1:] - ROOT_bones[:, :4]) / \
|
||||
(two_norm(ROOT_bones[:, 1:] - ROOT_bones[:, :4]) ** 2)
|
||||
PHI = angle_between(ROOT_bones[:, :4], ROOT_bones[:, 1:]) # (B)
|
||||
|
||||
rb_loss = interval_loss(value=curvatures, min=self.rb_curvatures_min, max=self.rb_curvatures_max) + \
|
||||
interval_loss(value=PHI, min=self.rb_PHI_min, max=self.rb_PHI_max)
|
||||
final_loss += self.lambda_rb * rb_loss
|
||||
BMC_losses["bmc_rb"] = rb_loss
|
||||
|
||||
# compute loss of Joint angles
|
||||
a_loss = torch.Tensor([0]).cuda()
|
||||
if self.lambda_a:
|
||||
# PIP bones
|
||||
PIP_X_axis = torch.zeros([batch_size, 5, 3]).cuda() # (B,5,3)
|
||||
PIP_X_axis[:, [0, 1, 4], :] = -normals[:, [0, 1, 3]]
|
||||
PIP_X_axis[:, 2:4] = -normalize(normals[:, 2:4] + normals[:, 1:3]) # (B,2,3)
|
||||
PIP_Y_axis = normalize(cross_product(PIP_Z_axis, PIP_X_axis)) # (B,5,3)
|
||||
|
||||
PIP_bones_xz = PIP_bones - inner_product(PIP_bones, PIP_Y_axis).unsqueeze(2) * PIP_Y_axis
|
||||
PIP_theta_flexion = angle_between(PIP_bones_xz, PIP_Z_axis) # in global coordinate (B)
|
||||
PIP_theta_abduction = angle_between(PIP_bones_xz, PIP_bones) # in global coordinate (B)
|
||||
# x-component of the bone vector
|
||||
tmp = inner_product(PIP_bones, PIP_X_axis)
|
||||
PIP_theta_flexion = torch.where(tmp < 1e-6, -PIP_theta_flexion, PIP_theta_flexion)
|
||||
# y-component of the bone vector
|
||||
tmp = torch.sum((PIP_bones * PIP_Y_axis), dim=-1)
|
||||
PIP_theta_abduction = torch.where(tmp < 1e-6, -PIP_theta_abduction, PIP_theta_abduction)
|
||||
|
||||
temp_axis = normalize(cross_product(PIP_Z_axis, PIP_bones))
|
||||
temp_alpha = angle_between(PIP_Z_axis, PIP_bones) # alpha belongs to [pi/2, pi]
|
||||
temp_R = axangle2mat_torch(axis=temp_axis, angle=temp_alpha, is_normalized=True)
|
||||
|
||||
# DIP bones
|
||||
DIP_X_axis = torch.matmul(temp_R, PIP_X_axis.unsqueeze(3)).squeeze()
|
||||
DIP_Y_axis = torch.matmul(temp_R, PIP_Y_axis.unsqueeze(3)).squeeze()
|
||||
|
||||
DIP_bones_xz = DIP_bones - inner_product(DIP_bones, DIP_Y_axis).unsqueeze(2) * DIP_Y_axis
|
||||
DIP_theta_flexion = angle_between(DIP_bones_xz, DIP_Z_axis) # in global coordinate
|
||||
DIP_theta_abduction = angle_between(DIP_bones_xz, DIP_bones) # in global coordinate
|
||||
# x-component of the bone vector
|
||||
tmp = inner_product(DIP_bones, DIP_X_axis)
|
||||
DIP_theta_flexion = torch.where(tmp < 0, -DIP_theta_flexion, DIP_theta_flexion)
|
||||
# y-component of the bone vector
|
||||
tmp = inner_product(DIP_bones, DIP_Y_axis)
|
||||
DIP_theta_abduction = torch.where(tmp < 0, -DIP_theta_abduction, DIP_theta_abduction)
|
||||
|
||||
temp_axis = normalize(cross_product(DIP_Z_axis, DIP_bones))
|
||||
temp_alpha = angle_between(DIP_Z_axis, DIP_bones) # alpha belongs to [pi/2, pi]
|
||||
temp_R = axangle2mat_torch(axis=temp_axis, angle=temp_alpha, is_normalized=True)
|
||||
|
||||
# TIP bones
|
||||
TIP_X_axis = torch.matmul(temp_R, DIP_X_axis.unsqueeze(3)).squeeze()
|
||||
TIP_Y_axis = torch.matmul(temp_R, DIP_Y_axis.unsqueeze(3)).squeeze()
|
||||
TIP_bones_xz = TIP_bones - inner_product(TIP_bones, TIP_Y_axis).unsqueeze(2) * TIP_Y_axis
|
||||
|
||||
TIP_theta_flexion = angle_between(TIP_bones_xz, TIP_Z_axis) # in global coordinate
|
||||
TIP_theta_abduction = angle_between(TIP_bones_xz, TIP_bones) # in global coordinate
|
||||
# x-component of the bone vector
|
||||
tmp = inner_product(TIP_bones, TIP_X_axis)
|
||||
TIP_theta_flexion = torch.where(tmp < 1e-6, -TIP_theta_flexion, TIP_theta_flexion)
|
||||
# y-component of the bone vector
|
||||
tmp = inner_product(TIP_bones, TIP_Y_axis)
|
||||
TIP_theta_abduction = torch.where(tmp < 1e-6, -TIP_theta_abduction, TIP_theta_abduction)
|
||||
|
||||
# ALL
|
||||
ALL_theta_flexion = torch.cat((PIP_theta_flexion, DIP_theta_flexion, TIP_theta_flexion), dim=-1)
|
||||
ALL_theta_abduction = torch.cat((PIP_theta_abduction, DIP_theta_abduction, TIP_theta_abduction), dim=-1)
|
||||
ALL_theta = torch.stack((ALL_theta_flexion, ALL_theta_abduction), dim=-1)
|
||||
|
||||
a_loss = calculate_joint_angle_loss(ALL_theta, self.joint_angle_limit)
|
||||
final_loss += self.lambda_a * a_loss
|
||||
|
||||
BMC_losses["bmc_a"] = a_loss
|
||||
|
||||
return final_loss, BMC_losses
|
||||
+180
@@ -0,0 +1,180 @@
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from body_model import OP_EDGE_LIST
|
||||
from geometry.rotation import batch_rodrigues
|
||||
from geometry.plane import (
|
||||
parse_floor_plane,
|
||||
compute_plane_intersection,
|
||||
fit_plane,
|
||||
)
|
||||
|
||||
|
||||
def compute_world2prior(floor_plane, trans, root_orient, origin):
|
||||
"""
|
||||
Computes rotation and translation from the camera frame to the canonical coordinate system
|
||||
used by the motion and initial state priors.
|
||||
- floor_plane : B x 3
|
||||
- trans : B x 3
|
||||
- root_orient : B x 3
|
||||
- origin: B x 3 desired origin (first joint of human model)
|
||||
- returns R (B, 3, 3), t (B, 3), root_height (B)
|
||||
"""
|
||||
B = trans.size(0)
|
||||
if floor_plane.size(1) == 3:
|
||||
floor_plane_4d = parse_floor_plane(floor_plane)
|
||||
else:
|
||||
floor_plane_4d = floor_plane
|
||||
floor_normal = floor_plane_4d[:, :3]
|
||||
floor_trans, _ = compute_plane_intersection(trans, -floor_normal, floor_plane_4d)
|
||||
|
||||
# compute prior frame axes within the camera frame
|
||||
# up is the floor_plane normal
|
||||
up_axis = floor_normal
|
||||
# right is body -x direction projected to floor plane
|
||||
root_orient_mat = batch_rodrigues(root_orient)
|
||||
body_right = -root_orient_mat[:, :, 0]
|
||||
floor_body_right, s = compute_plane_intersection(trans, body_right, floor_plane_4d)
|
||||
right_axis = floor_body_right - floor_trans
|
||||
# body right may not actually intersect - in this case must negate axis because we have the -x
|
||||
right_axis = torch.where(s.reshape((B, 1)) < 0, -right_axis, right_axis)
|
||||
right_axis = right_axis / torch.norm(right_axis, dim=1, keepdim=True)
|
||||
# forward is their cross product
|
||||
fwd_axis = torch.linalg.cross(up_axis, right_axis)
|
||||
fwd_axis = fwd_axis / torch.norm(fwd_axis, dim=1, keepdim=True)
|
||||
|
||||
# prior frame is right, fwd, up
|
||||
prior_R = torch.stack([right_axis, fwd_axis, up_axis], dim=2)
|
||||
world2prior_R = prior_R.transpose(-1, -2)
|
||||
|
||||
# translation takes translation to origin plus offset to the floor
|
||||
world2prior_t = -trans
|
||||
|
||||
# compute the distance from root to ground plane
|
||||
_, s_root = compute_plane_intersection(origin, -floor_normal, floor_plane_4d)
|
||||
root_height = s_root.reshape(B, 1)
|
||||
|
||||
if False:
|
||||
# apply transform to origin
|
||||
origin_prior = torch.matmul(world2prior_R, (origin - trans)[..., None])[..., 0]
|
||||
root_height = s_root.reshape(B, 1) - origin_prior[:, 2:3]
|
||||
|
||||
return world2prior_R, world2prior_t, root_height
|
||||
|
||||
|
||||
def find_cliques(edges):
|
||||
cliques = []
|
||||
N = edges.shape[0]
|
||||
empty = np.ones(N, dtype=bool)
|
||||
while empty.sum() > 0:
|
||||
idcs = np.where(empty)[0]
|
||||
i = idcs[0]
|
||||
cur_clique = [i]
|
||||
for j in idcs[1:]:
|
||||
if all(edges[c, j] for c in cur_clique):
|
||||
cur_clique.append(j)
|
||||
empty[cur_clique] = False
|
||||
cliques.append(cur_clique)
|
||||
return cliques
|
||||
|
||||
|
||||
def estimate_floor_planes(
|
||||
smpl_joints, valid_mask, thresh=0.5, group=False, flatten=False
|
||||
):
|
||||
"""
|
||||
:param smpl_joints (B, T, J, 3)
|
||||
:param valid_mask (B, T)
|
||||
:param thresh (optional, default 0.5) threshold for separating
|
||||
:param group (optional, default False) whether to group the planes
|
||||
:param flatten (optional, default False) whether to estimate one plane for everyone
|
||||
returns floor_plane (C, 3), floor_idcs (B,)
|
||||
"""
|
||||
B, T = valid_mask.shape
|
||||
device = valid_mask.device
|
||||
feet_joints = smpl_joints[..., 10:12, :] # (B, T, 2, 3)
|
||||
points = [feet_joints[b, valid_mask[b]].reshape(-1, 3) for b in range(B)]
|
||||
|
||||
if flatten:
|
||||
planes = fit_plane(torch.cat(points, dim=0))[None]
|
||||
labels = torch.zeros(B, device=device, dtype=torch.long)
|
||||
return planes[..., :3] * planes[..., 3:], labels
|
||||
|
||||
# estimate floors independently
|
||||
planes = torch.stack([fit_plane(pts) for pts in points], dim=0) # (B, 4)
|
||||
normals, offsets = planes[..., :3], planes[..., 3:] # (B, 3), (B, 1)
|
||||
sgn = torch.sign(offsets)
|
||||
normals, offsets = sgn * normals, sgn * offsets
|
||||
planes = normals * offsets # (B, 3)
|
||||
labels = torch.arange(B, device=device)
|
||||
|
||||
if not group:
|
||||
return planes, labels # (B, 3), (B,)
|
||||
|
||||
# asymetric, project each plane vector onto the other normals
|
||||
diffs = offsets - (planes[None] * normals[:, None]).sum(dim=-1)
|
||||
print(diffs)
|
||||
edges = diffs < thresh
|
||||
edges = edges.T & edges
|
||||
clusters = find_cliques(edges) # list of indices of clusters
|
||||
planes = []
|
||||
labels = torch.zeros(B, dtype=torch.long, device=device)
|
||||
for c, idcs in enumerate(clusters):
|
||||
plane = fit_plane(torch.cat([points[i] for i in idcs], dim=0))
|
||||
planes.append(plane)
|
||||
labels[idcs] = c
|
||||
planes = torch.stack(planes, dim=0) # (C, 4)
|
||||
return planes[..., :3] * planes[..., 3:], labels
|
||||
|
||||
|
||||
def estimate_initial_trans(body_pose, joints3d_op, joints2d_op, focal):
|
||||
"""
|
||||
use focal length and bone lengths to approximate distance from camera
|
||||
(based on PROX https://github.com/mohamedhassanmus/prox/blob/master/prox/fitting.py)
|
||||
:param body_pose
|
||||
:param joints3d_op OP keypoints in 3d
|
||||
:param joints2d_op OP keypoints in 2d
|
||||
"""
|
||||
B, T, N, _ = joints2d_op.shape
|
||||
joints2d_obs = joints2d_op[:, :, :, :2] # (B, T, N, 3)
|
||||
joints2d_conf = joints2d_op[:, :, :, 2]
|
||||
|
||||
# find least-occluded 2d frame
|
||||
num_2d_vis = torch.sum(joints2d_conf > 0.0, dim=2)
|
||||
best_2d_idx = torch.max(num_2d_vis, dim=1)[1]
|
||||
|
||||
# calculate bone lengths and confidence in each bone length
|
||||
bone3d = []
|
||||
bone2d = []
|
||||
conf2d = []
|
||||
for pair in OP_EDGE_LIST:
|
||||
diff3d = torch.norm(
|
||||
joints3d_op[:, 0, pair[0], :] - joints3d_op[:, 0, pair[1], :],
|
||||
dim=1,
|
||||
) # does not change over time
|
||||
diff2d = torch.norm(
|
||||
joints2d_obs[:, :, pair[0], :] - joints2d_obs[:, :, pair[1], :],
|
||||
dim=2,
|
||||
)
|
||||
minconf2d = torch.min(
|
||||
joints2d_conf[:, :, pair[0]], joints2d_conf[:, :, pair[1]]
|
||||
)
|
||||
bone3d.append(diff3d)
|
||||
bone2d.append(diff2d)
|
||||
conf2d.append(minconf2d)
|
||||
|
||||
bone3d = torch.stack(bone3d, dim=1)
|
||||
bone2d = torch.stack(bone2d, dim=2)
|
||||
bone2d = bone2d[np.arange(B), best_2d_idx, :]
|
||||
conf2d = torch.stack(conf2d, dim=2)
|
||||
conf2d = conf2d[np.arange(B), best_2d_idx, :]
|
||||
|
||||
# mean over all bones
|
||||
mean_bone3d = torch.mean(bone3d, dim=1)
|
||||
mean_bone2d = torch.mean(bone2d * (conf2d > 0.0), dim=1)
|
||||
|
||||
# approx z based on ratio
|
||||
init_z = focal * (mean_bone3d / mean_bone2d)
|
||||
init_trans = torch.zeros(B, T, 3)
|
||||
init_trans[:, :, 2] = init_z.unsqueeze(1).expand(B, T).detach()
|
||||
|
||||
return init_trans
|
||||
+1456
File diff suppressed because it is too large
Load Diff
+727
@@ -0,0 +1,727 @@
|
||||
import os
|
||||
import numpy as np
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from body_model import MANO_JOINTS, KEYPT_VERTS, smpl_to_openpose
|
||||
from geometry.rotation import batch_rodrigues, rotation_matrix_to_angle_axis
|
||||
from util.logger import Logger
|
||||
from util.tensor import (
|
||||
detach_all,
|
||||
scatter_intervals,
|
||||
select_intervals,
|
||||
get_scatter_mask,
|
||||
)
|
||||
|
||||
from .base_scene import BaseSceneModel
|
||||
from .helpers import compute_world2prior, estimate_floor_planes
|
||||
|
||||
|
||||
J_BODY = len(MANO_JOINTS) - 1 # no root
|
||||
CONTACT_ORDERING = [
|
||||
"hips",
|
||||
"leftLeg",
|
||||
"rightLeg",
|
||||
"leftFoot",
|
||||
"rightFoot",
|
||||
"leftToeBase",
|
||||
"rightToeBase",
|
||||
"leftHand",
|
||||
"rightHand",
|
||||
]
|
||||
CONTACT_INDS = [MANO_JOINTS[jname] for jname in CONTACT_ORDERING]
|
||||
CONTACT_THRESH = 0.5
|
||||
|
||||
|
||||
class MovingSceneModel(BaseSceneModel):
|
||||
"""
|
||||
Scene model of moving people in a shared global reference frame
|
||||
|
||||
Parameters:
|
||||
batch_size: number of sequences to optimize
|
||||
seq_len: length of the sequences
|
||||
pose_prior: VPoser model
|
||||
motion_prior: humor model
|
||||
init_motion_prior: dict of GMM params to use for prior on initial motion state
|
||||
shared_floor: (default True) if true, sequences are in the same reference frame
|
||||
fit_gender: (optional) gender of SMPL model
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
batch_size,
|
||||
seq_len,
|
||||
body_model,
|
||||
pose_prior,
|
||||
motion_prior,
|
||||
init_motion_prior=None,
|
||||
fit_gender="neutral",
|
||||
use_init=False,
|
||||
opt_cams=False,
|
||||
opt_scale=True,
|
||||
cam_graph=False,
|
||||
est_floor=True,
|
||||
floor_type="shared",
|
||||
async_tracks=True,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__(
|
||||
batch_size,
|
||||
seq_len,
|
||||
body_model,
|
||||
pose_prior,
|
||||
fit_gender=fit_gender,
|
||||
use_init=use_init,
|
||||
opt_cams=opt_cams,
|
||||
opt_scale=opt_scale,
|
||||
cam_graph=cam_graph,
|
||||
)
|
||||
assert motion_prior is not None
|
||||
assert motion_prior.model_data_config in [
|
||||
"mano+joints",
|
||||
"mano+joints+contacts",
|
||||
], "Only mano+joints motion prior configuration is supported!"
|
||||
|
||||
self.motion_prior = motion_prior
|
||||
self.init_motion_prior = init_motion_prior
|
||||
|
||||
# need latent dynamics sequence as well
|
||||
self.latent_motion_dim = self.motion_prior.latent_size
|
||||
self.cond_prior = self.motion_prior.use_conditional_prior
|
||||
assert self.cond_prior
|
||||
|
||||
# the frame chosen to use for the initial state (first frame by default)
|
||||
self.async_tracks = async_tracks
|
||||
self.register_buffer("track_start", torch.zeros(self.batch_size))
|
||||
self.register_buffer("track_end", torch.ones(self.batch_size) * self.seq_len)
|
||||
|
||||
self.shared_floor = floor_type == "shared"
|
||||
self.group_floor = floor_type == "group"
|
||||
self.floor_type = floor_type
|
||||
print("FLOOR TYPE", floor_type)
|
||||
self.est_floor = est_floor
|
||||
|
||||
@property
|
||||
def is_motion_active(self):
|
||||
return hasattr(self.params, "latent_motion")
|
||||
|
||||
def initialize(self, obs_data, cam_data, param_dict, data_fps):
|
||||
"""
|
||||
we need to also optimize for floor and world scale
|
||||
"""
|
||||
Logger.log("Initializing moving scene model with observed data")
|
||||
|
||||
self.params.set_cameras(
|
||||
cam_data,
|
||||
opt_scale=self.opt_scale,
|
||||
opt_cams=self.opt_cams,
|
||||
opt_focal=self.opt_cams,
|
||||
**param_dict,
|
||||
)
|
||||
# self.init_floor(obs_data, param_dict)
|
||||
self.init_first_state(obs_data, param_dict, data_fps)
|
||||
|
||||
# def init_floor(self, obs_data, param_dict):
|
||||
# if self.est_floor or self.group_floor:
|
||||
# with torch.no_grad():
|
||||
# smpl_preds = self.pred_mano(
|
||||
# param_dict["trans"],
|
||||
# param_dict["root_orient"],
|
||||
# self.latent2pose(param_dict["latent_pose"]),
|
||||
# param_dict["betas"],
|
||||
# )
|
||||
# floor_plane, floor_idcs = estimate_floor_planes(
|
||||
# smpl_preds["joints3d"].detach(),
|
||||
# obs_data["vis_mask"] > 0,
|
||||
# group=self.group_floor,
|
||||
# flatten=self.shared_floor,
|
||||
# )
|
||||
# if self.group_floor and not self.est_floor:
|
||||
# # don't use the estimated floor as initial plane
|
||||
# floor_plane = obs_data["floor_plane"][: len(floor_plane)]
|
||||
# else: # fixed shared or separate floors
|
||||
# num_floors = 1 if self.shared_floor else self.batch_size
|
||||
# floor_plane = obs_data["floor_plane"][:num_floors]
|
||||
# if self.shared_floor:
|
||||
# floor_idcs = torch.zeros(
|
||||
# self.batch_size, dtype=torch.long, device=floor_plane.device
|
||||
# )
|
||||
# else:
|
||||
# floor_idcs = torch.arange(
|
||||
# self.batch_size, dtype=torch.long, device=floor_plane.device
|
||||
# )
|
||||
|
||||
# Logger.log(f"ESTIMATED FLOORS: {str(floor_plane.detach().cpu())}")
|
||||
# Logger.log(f"FLOOR IDCS: {str(floor_idcs.detach().cpu())}")
|
||||
# self.params.set_param("floor_plane", floor_plane.float().detach())
|
||||
# self.params.set_param(
|
||||
# "floor_idcs", floor_idcs.long().detach(), requires_grad=False
|
||||
# )
|
||||
|
||||
def init_first_state(self, obs_data, param_dict, data_fps):
|
||||
"""
|
||||
initialize the latent motion and first state of each track
|
||||
using per-frame trajectories of trans, rot, latent_pose, betas
|
||||
and observed data
|
||||
"""
|
||||
B, T = self.batch_size, self.seq_len
|
||||
param_dict = detach_all(param_dict)
|
||||
|
||||
# select the valid segments of each track; pad the shorter tracks with final element
|
||||
if self.async_tracks and "track_interval" in obs_data:
|
||||
print("asynchronous tracks")
|
||||
interval = obs_data["track_interval"] # (B, 2)
|
||||
start, end = interval[:, 0], interval[:, 1]
|
||||
self.track_start, self.track_end = start, end
|
||||
param_dict = select_dict_segments(
|
||||
param_dict, start, end, names=["trans", "root_orient", "latent_pose"]
|
||||
)
|
||||
|
||||
trans = param_dict["trans"] # (B, T, 3)
|
||||
root_orient = param_dict["root_orient"] # (B, T, 3)
|
||||
latent_pose = param_dict["latent_pose"] # (B, T, D)
|
||||
betas = param_dict["betas"] # (B, b)
|
||||
|
||||
Logger.log(f"SEL TRACKS {trans.shape}, {root_orient.shape}")
|
||||
|
||||
# save each track's first appearance
|
||||
self.params.set_param("trans", trans[:, :1])
|
||||
Logger.log(f"INITIAL TRANS {trans[:, :1].detach().cpu()}")
|
||||
self.params.set_param("root_orient", root_orient[:, :1])
|
||||
self.params.set_param("latent_pose", latent_pose[:, :1])
|
||||
self.params.set_param("betas", betas)
|
||||
|
||||
# pass the current pose estimates through the motion prior
|
||||
self.data_fps = data_fps
|
||||
body_pose = self.latent2pose(latent_pose)
|
||||
init_latent = self.infer_latent_motion(
|
||||
trans, root_orient, body_pose, betas, data_fps
|
||||
).detach() # (B, T-1, D)
|
||||
self.params.set_param("latent_motion", init_latent)
|
||||
|
||||
# estimate velocities in prior frame and save initial state
|
||||
trans_vel, joints_vel, root_orient_vel = self.estimate_prior_velocities(
|
||||
trans, root_orient, body_pose, betas, data_fps
|
||||
)
|
||||
self.params.set_param("trans_vel", trans_vel[:, :1].detach())
|
||||
self.params.set_param("joints_vel", joints_vel[:, :1].detach())
|
||||
self.params.set_param("root_orient_vel", root_orient_vel[:, :1].detach())
|
||||
|
||||
def get_optim_result(self, num_steps=-1):
|
||||
res = super().get_optim_result()
|
||||
if not self.is_motion_active:
|
||||
return res
|
||||
|
||||
num_steps = self.seq_len if num_steps < 0 else num_steps
|
||||
prior_res, world_res = self.rollout_latent_motion(
|
||||
self.params.latent_motion[:, : num_steps - 1]
|
||||
)
|
||||
world_res = self.synchronize_preds(world_res, num_steps)
|
||||
res["world"].update(world_res)
|
||||
|
||||
prior_res = self.synchronize_preds(prior_res, num_steps)
|
||||
res["prior"] = prior_res
|
||||
return res
|
||||
|
||||
def estimate_prior_velocities(self, trans, root_orient, body_pose, betas, data_fps):
|
||||
"""
|
||||
compute velocities of trajectory in prior frame (not world frame)
|
||||
Transforms the first of each track into the prior frame
|
||||
:param trans (B, T, 3)
|
||||
:param root_orient (B, T, 3)
|
||||
:param body_pose (B, T, Dp)
|
||||
:param betas (B, Db)
|
||||
"""
|
||||
with torch.no_grad():
|
||||
self.update_world2prior(trans, root_orient, body_pose, betas)
|
||||
trans, root_orient = self.apply_world2prior(
|
||||
trans, root_orient, body_pose, betas
|
||||
)
|
||||
mano_results = self.pred_mano(trans, root_orient, body_pose, betas)
|
||||
return estimate_velocities(
|
||||
trans, root_orient, mano_results["joints3d"], data_fps
|
||||
)
|
||||
|
||||
def update_world2prior(self, trans, root_orient, body_pose, betas):
|
||||
B = trans.shape[0]
|
||||
mano_data = self.pred_mano(
|
||||
trans[:, :1], root_orient[:, :1], body_pose[:, :1], betas
|
||||
)
|
||||
floor_plane = self.params.floor_plane[self.params.floor_idcs]
|
||||
R, t, height = compute_world2prior(
|
||||
floor_plane,
|
||||
trans[:, 0],
|
||||
root_orient[:, 0],
|
||||
mano_data["joints3d"][:, 0, 0],
|
||||
)
|
||||
self.world2prior_R = R # (B, 3, 3)
|
||||
self.world2prior_t = t # (B, 3)
|
||||
self.world2prior_root_height = height # (B)
|
||||
|
||||
def apply_world2prior(self, trans, root_orient, body_pose, betas, inverse=False):
|
||||
"""
|
||||
Applies the world2prior tranformation to trans, root_orient
|
||||
If T=1, this function assumes they are at key_frame_idx,
|
||||
which we need to compute the offset from the origin
|
||||
:param trans (B, T, 3)
|
||||
:param root_orient (B, T, 3)
|
||||
:param body_pose (B, T, J, 3)
|
||||
:param betas (B, b)
|
||||
:param inverse (bool) optional, default False
|
||||
"""
|
||||
B, T, _ = trans.size()
|
||||
# (B, 3, 3), (B, 3), (B)
|
||||
R, t, root_height = (
|
||||
self.world2prior_R,
|
||||
self.world2prior_t,
|
||||
self.world2prior_root_height,
|
||||
)
|
||||
R_time = R.unsqueeze(1).expand((B, T, 3, 3))
|
||||
t_time = t.unsqueeze(1).expand((B, T, 3))
|
||||
root_orient_mat = batch_rodrigues(root_orient.reshape(-1, 3)).reshape(
|
||||
B, T, 3, 3
|
||||
)
|
||||
|
||||
if inverse:
|
||||
R_time = R_time.transpose(-1, -2)
|
||||
|
||||
root_orient_mat = torch.matmul(R_time, root_orient_mat)
|
||||
root_orient = rotation_matrix_to_angle_axis(
|
||||
root_orient_mat.reshape(B * T, 3, 3)
|
||||
).reshape(B, T, 3)
|
||||
|
||||
if inverse:
|
||||
# transform so first frame is at origin
|
||||
trans = trans - trans[:, :1, :]
|
||||
|
||||
# rotates to world frame
|
||||
trans = torch.matmul(R_time, trans[..., None])[..., 0]
|
||||
# translate to world frame
|
||||
trans = trans - t_time
|
||||
|
||||
return trans, root_orient
|
||||
|
||||
# first transform so the trans of key frame is at origin
|
||||
trans = trans + t_time
|
||||
# then rotate to canonical frame
|
||||
trans = torch.matmul(R_time, trans[..., None])[..., 0]
|
||||
# compute the root height after transforming
|
||||
cur_smpl_data = self.pred_mano(trans, root_orient, body_pose, betas)
|
||||
cur_root_height = cur_smpl_data["joints3d"][:, 0, 0, 2:3]
|
||||
# then apply floor offset so the root joint is at root_height
|
||||
height_diff = root_height - cur_root_height
|
||||
trans_offset = torch.cat(
|
||||
[torch.zeros((B, 2)).to(height_diff), height_diff], axis=1
|
||||
)
|
||||
trans = trans + trans_offset.reshape((B, 1, 3))
|
||||
|
||||
return trans, root_orient
|
||||
|
||||
def convert_prior_rot_inputs(self, root_orient, body_pose):
|
||||
# convert rots
|
||||
# body pose and root orient are both in aa
|
||||
B, T = root_orient.shape[:2]
|
||||
if (
|
||||
self.motion_prior.in_rot_rep == "mat"
|
||||
or self.motion_prior.in_rot_rep == "6d"
|
||||
):
|
||||
root_orient_in = batch_rodrigues(root_orient.reshape(-1, 3)).reshape(
|
||||
(B, T, 9)
|
||||
)
|
||||
body_pose_in = batch_rodrigues(body_pose.reshape(-1, 3)).reshape(
|
||||
(B, T, J_BODY * 9)
|
||||
)
|
||||
if self.motion_prior.in_rot_rep == "6d":
|
||||
root_orient_in = root_orient[:, :, :6]
|
||||
body_pose_in = body_pose.reshape((B, T, J_BODY, 9))[:, :, :, :6].reshape(
|
||||
(B, T, J_BODY * 6)
|
||||
)
|
||||
return root_orient_in, body_pose_in
|
||||
|
||||
def convert_prior_rot_outputs(self, out_dict):
|
||||
keys = ["root_orient", "pose_body"]
|
||||
for k in keys:
|
||||
out = out_dict[k]
|
||||
B, T = out.shape[:2]
|
||||
out_dict[k] = rotation_matrix_to_angle_axis(
|
||||
out.reshape((-1, 3, 3))
|
||||
).reshape((B, T, -1))
|
||||
return out_dict
|
||||
|
||||
def infer_latent_motion(
|
||||
self, trans, root_orient, body_pose, betas, data_fps, full_forward_pass=False
|
||||
):
|
||||
"""
|
||||
By default, gets a sequence of z's from the current SMPL optim params.
|
||||
|
||||
If full_forward_pass is true, also samples from the posterior and feeds
|
||||
through the motion prior decoder to get all terms needed to calculate the ELBO.
|
||||
"""
|
||||
B, T, _ = trans.size()
|
||||
h = 1.0 / data_fps
|
||||
|
||||
# need to first transform into canonical coordinate frame
|
||||
self.update_world2prior(trans, root_orient, body_pose, betas)
|
||||
trans, root_orient = self.apply_world2prior(
|
||||
trans, root_orient, body_pose, betas
|
||||
)
|
||||
|
||||
mano_results = self.pred_mano(trans, root_orient, body_pose, betas)
|
||||
joints = mano_results["joints3d"] # (B, T, len(MANO_JOINTS), 3)
|
||||
trans_vel, joints_vel, root_orient_vel = estimate_velocities(
|
||||
trans, root_orient, joints, data_fps
|
||||
)
|
||||
|
||||
root_orient_in, body_pose_in = self.convert_prior_rot_inputs(
|
||||
root_orient, body_pose
|
||||
)
|
||||
joints_in = joints.reshape((B, T, -1))
|
||||
joints_vel_in = joints_vel.reshape((B, T, -1))
|
||||
|
||||
seq_dict = {
|
||||
"trans": trans,
|
||||
"trans_vel": trans_vel,
|
||||
"root_orient": root_orient_in,
|
||||
"root_orient_vel": root_orient_vel,
|
||||
"pose_body": body_pose_in,
|
||||
"joints": joints_in,
|
||||
"joints_vel": joints_vel_in,
|
||||
}
|
||||
|
||||
infer_results = self.motion_prior.infer_global_seq(
|
||||
seq_dict, full_forward_pass=full_forward_pass
|
||||
)
|
||||
if full_forward_pass:
|
||||
# return both the given motion and the one from the forward pass
|
||||
# make sure rotations are matrix
|
||||
# NOTE: assumes seq_dict is same thing we want to compute loss on
|
||||
# Need to change if multiple future steps.
|
||||
if self.motion_prior.in_rot_rep != "mat":
|
||||
seq_dict["trans"] = batch_rodrigues(root_orient.reshape(-1, 3)).reshape(
|
||||
(B, T, 9)
|
||||
)
|
||||
seq_dict["pose_body"] = batch_rodrigues(
|
||||
body_pose.reshape(-1, 3)
|
||||
).reshape((B, T, J_BODY * 9))
|
||||
# do not need initial step anymore since output will be T-1
|
||||
for k, v in seq_dict.items():
|
||||
seq_dict[k] = v[:, 1:]
|
||||
for k in infer_results.keys():
|
||||
if k != "posterior_distrib" and k != "prior_distrib":
|
||||
infer_results[k] = infer_results[k][
|
||||
:, :, 0
|
||||
] # only want first output step
|
||||
infer_results = (seq_dict, infer_results)
|
||||
else:
|
||||
prior_z, posterior_z = infer_results
|
||||
infer_results = posterior_z[0] # mean of the approximate posterior
|
||||
|
||||
return infer_results
|
||||
|
||||
def rollout_mano_steps(self, num_steps=-1):
|
||||
# rollout motion given initial state
|
||||
num_steps = self.seq_len if num_steps < 0 else num_steps
|
||||
latent_motion = self.params.latent_motion[:, : num_steps - 1]
|
||||
|
||||
# roll out tracks from their first appearances
|
||||
# NOTE rolled out tracks are unsynced in time
|
||||
res, world_res = self.rollout_latent_motion(
|
||||
latent_motion, return_prior=self.cond_prior
|
||||
)
|
||||
|
||||
# get the smpl predictions for the unsynced tracks
|
||||
preds = self.pred_mano(
|
||||
res["trans"], res["root_orient"], res["pose_body"], res["betas"]
|
||||
)
|
||||
world_preds = self.pred_mano(
|
||||
world_res["trans"],
|
||||
world_res["root_orient"],
|
||||
world_res["pose_body"],
|
||||
world_res["betas"],
|
||||
)
|
||||
|
||||
# pass along relevant results
|
||||
preds["joints3d_rollout"] = res["joints"]
|
||||
preds["latent_pose"] = self.pose2latent(res["pose_body"])
|
||||
preds["joints3d_init"] = preds["joints3d"][:, :1]
|
||||
|
||||
if "contacts" in res:
|
||||
preds["contacts"] = res["contacts"]
|
||||
if "contacts_conf" in res:
|
||||
preds["contacts_conf"] = res["contacts_conf"]
|
||||
|
||||
# synchronize the tracklets to the same timesteps
|
||||
preds = self.synchronize_preds(preds, num_steps)
|
||||
world_preds = self.synchronize_preds(world_preds, num_steps)
|
||||
|
||||
# pass along unsynchronized motion latents
|
||||
preds["latent_motion"] = latent_motion
|
||||
if "cond_prior" in res:
|
||||
preds["cond_prior"] = res["cond_prior"]
|
||||
|
||||
return preds, world_preds
|
||||
|
||||
def rollout_latent_motion(
|
||||
self,
|
||||
latent_motion,
|
||||
return_prior=False,
|
||||
return_vel=False,
|
||||
num_steps=-1,
|
||||
canonicalize_input=False,
|
||||
):
|
||||
"""
|
||||
From the stored initial state, rolls out a sequence from the latent_motion vector
|
||||
NOTE: the stored initial states are not at the same time step, but we return all
|
||||
predicted sequences starting from time 0.
|
||||
|
||||
If latent_motion is None, samples num_steps into the future sequence from the prior.
|
||||
using the mean of the prior rather than random samples.
|
||||
|
||||
If canonicalize_input is True, transform the initial state into the local canonical
|
||||
frame before roll out
|
||||
"""
|
||||
# get the first frame state
|
||||
# NOTE: first states do not occur at same time step
|
||||
trans, root_orient, betas = (
|
||||
self.params.trans,
|
||||
self.params.root_orient,
|
||||
self.params.betas,
|
||||
)
|
||||
trans_vel, root_orient_vel, joints_vel = (
|
||||
self.params.trans_vel,
|
||||
self.params.root_orient_vel,
|
||||
self.params.joints_vel,
|
||||
)
|
||||
body_pose = self.latent2pose(self.params.latent_pose)
|
||||
|
||||
B = trans.size(0)
|
||||
is_sampling = latent_motion is None
|
||||
Tm1 = num_steps if latent_motion is None else latent_motion.size(1)
|
||||
if is_sampling and Tm1 <= 0:
|
||||
Logger.log("num_steps must be positive to sample!")
|
||||
exit()
|
||||
|
||||
# need to first transform initial state into canonical coordinate frame
|
||||
self.update_world2prior(trans, root_orient, body_pose, betas)
|
||||
trans, root_orient = self.apply_world2prior(
|
||||
trans, root_orient, body_pose, betas
|
||||
)
|
||||
|
||||
smpl_results = self.pred_mano(trans, root_orient, body_pose, betas)
|
||||
joints = smpl_results["joints3d"] # (B, T, len(MANO_JOINTS), 3)
|
||||
|
||||
# update to correct rotations for input
|
||||
root_orient_in, body_pose_in = self.convert_prior_rot_inputs(
|
||||
root_orient, body_pose
|
||||
)
|
||||
joints_in = joints.reshape((B, 1, -1))
|
||||
joints_vel_in = joints_vel.reshape((B, 1, -1))
|
||||
|
||||
rollout_in_dict = {
|
||||
"trans": trans,
|
||||
"root_orient": root_orient_in,
|
||||
"pose_body": body_pose_in,
|
||||
"joints": joints_in,
|
||||
"trans_vel": trans_vel,
|
||||
"root_orient_vel": root_orient_vel,
|
||||
"joints_vel": joints_vel_in,
|
||||
}
|
||||
|
||||
roll_output = self.motion_prior.roll_out(
|
||||
None,
|
||||
rollout_in_dict,
|
||||
Tm1,
|
||||
z_seq=latent_motion,
|
||||
return_prior=return_prior,
|
||||
return_z=is_sampling,
|
||||
canonicalize_input=canonicalize_input,
|
||||
gender=[self.fit_gender] * B,
|
||||
betas=betas.reshape((B, 1, -1)),
|
||||
)
|
||||
|
||||
pred_dict = prior_out = None
|
||||
if return_prior:
|
||||
pred_dict, prior_out = roll_output
|
||||
else:
|
||||
pred_dict = roll_output
|
||||
|
||||
pred_dict = self.convert_prior_rot_outputs(pred_dict)
|
||||
|
||||
# concat with initial state
|
||||
trans_out = torch.cat([trans, pred_dict["trans"]], dim=1)
|
||||
root_orient_out = torch.cat([root_orient, pred_dict["root_orient"]], dim=1)
|
||||
body_pose_out = torch.cat([body_pose, pred_dict["pose_body"]], dim=1)
|
||||
joints_out = torch.cat(
|
||||
[joints, pred_dict["joints"].reshape((B, Tm1, -1, 3))], dim=1
|
||||
)
|
||||
out_dict = {
|
||||
"trans": trans_out,
|
||||
"root_orient": root_orient_out,
|
||||
"pose_body": body_pose_out,
|
||||
"joints": joints_out,
|
||||
"betas": betas,
|
||||
}
|
||||
if return_vel:
|
||||
out_dict["trans_vel"] = torch.cat(
|
||||
[trans_vel, pred_dict["trans_vel"]], dim=1
|
||||
)
|
||||
out_dict["root_orient_vel"] = torch.cat(
|
||||
[root_orient_vel, pred_dict["root_orient_vel"]], dim=1
|
||||
)
|
||||
out_dict["joints_vel"] = torch.cat(
|
||||
[joints_vel, pred_dict["joints_vel"].reshape((B, Tm1, -1, 3))],
|
||||
dim=1,
|
||||
)
|
||||
|
||||
if return_prior: # return the mean and var of distribution stacked
|
||||
pm, pv = prior_out
|
||||
out_dict["cond_prior"] = torch.stack([pm, pv], dim=-1)
|
||||
|
||||
if self.motion_prior.model_data_config == "smpl+joints+contacts":
|
||||
pred_contacts = pred_dict["contacts"]
|
||||
# get binary classification
|
||||
contact_conf = torch.sigmoid(pred_contacts)
|
||||
pred_contacts = (contact_conf > CONTACT_THRESH).to(torch.float)
|
||||
# expand to full body
|
||||
full_contact_conf = torch.zeros((B, Tm1, len(MANO_JOINTS))).to(contact_conf)
|
||||
full_contact_conf[:, :, CONTACT_INDS] = (
|
||||
full_contact_conf[:, :, CONTACT_INDS] + contact_conf
|
||||
)
|
||||
full_contacts = torch.zeros((B, Tm1, len(MANO_JOINTS))).to(pred_contacts)
|
||||
full_contacts[:, :, CONTACT_INDS] = (
|
||||
full_contacts[:, :, CONTACT_INDS] + pred_contacts
|
||||
)
|
||||
# repeat first entry for t0
|
||||
full_contact_conf = torch.cat(
|
||||
[full_contact_conf[:, 0:1], full_contact_conf], dim=1
|
||||
)
|
||||
full_contacts = torch.cat([full_contacts[:, 0:1], full_contacts], dim=1)
|
||||
out_dict["contacts_conf"] = full_contact_conf
|
||||
out_dict["contacts"] = full_contacts
|
||||
|
||||
if is_sampling:
|
||||
out_dict["z"] = pred_dict["z"]
|
||||
|
||||
cam_dict = {
|
||||
"trans": out_dict["trans"],
|
||||
"root_orient": out_dict["root_orient"],
|
||||
"pose_body": out_dict["pose_body"],
|
||||
"betas": out_dict["betas"],
|
||||
}
|
||||
|
||||
# also must return trans and root orient in camera frame
|
||||
cam_dict["trans"], cam_dict["root_orient"] = self.apply_world2prior(
|
||||
out_dict["trans"],
|
||||
out_dict["root_orient"],
|
||||
out_dict["pose_body"],
|
||||
out_dict["betas"],
|
||||
inverse=True,
|
||||
)
|
||||
return out_dict, cam_dict
|
||||
|
||||
def synchronize_preds(self, pred_dict, seg_len):
|
||||
"""
|
||||
synchronize predictions in time, scatter predictions into [0, seq_len)
|
||||
"""
|
||||
if not self.async_tracks:
|
||||
# predictions are already synchronized
|
||||
return pred_dict, None
|
||||
|
||||
# return time-synchronized predictions
|
||||
start = self.track_start
|
||||
end = torch.clip(start + seg_len, max=self.track_end)
|
||||
T = end.max()
|
||||
pred_dict = scatter_dict_segments(pred_dict, start, end, T)
|
||||
pred_mask = get_scatter_mask(start, end, T)
|
||||
pred_dict["track_mask"] = pred_mask
|
||||
return pred_dict
|
||||
|
||||
|
||||
def scatter_dict_segments(data_dict, start, end, T=None, names=None):
|
||||
"""
|
||||
the tracks as they are stored are synchronized by time step
|
||||
Uses the start and end to synchronize output predictions
|
||||
"""
|
||||
sync_dict = data_dict.copy()
|
||||
min_len = (end - start).min()
|
||||
if names is None:
|
||||
names = data_dict.keys()
|
||||
for name in names:
|
||||
val = data_dict[name]
|
||||
if not isinstance(val, torch.Tensor) or val.ndim < 3 or val.shape[1] < min_len:
|
||||
continue
|
||||
sync_dict[name] = scatter_intervals(val, start, end, T)
|
||||
return sync_dict
|
||||
|
||||
|
||||
def select_dict_segments(data_dict, start, end, names=None):
|
||||
out_data_dict = data_dict.copy()
|
||||
min_len = (end - start).min()
|
||||
if names is None:
|
||||
names = data_dict.keys()
|
||||
for name in names:
|
||||
val = data_dict[name]
|
||||
if not isinstance(val, torch.Tensor) or val.ndim < 3 or val.shape[1] < min_len:
|
||||
continue
|
||||
# only select for time-series observations
|
||||
out_data_dict[name] = select_intervals(val, start, end)
|
||||
return out_data_dict
|
||||
|
||||
|
||||
def estimate_velocities(trans, root_orient, joints3d, data_fps):
|
||||
"""
|
||||
Estimates velocity inputs to the motion prior.
|
||||
- trans (B, T, 3) root translation
|
||||
- root_orient (B, T, 3) aa root orientation
|
||||
- joints3d (B, T, len(MANO_JOINTS), 3) joints3d of SMPL prediction
|
||||
"""
|
||||
B, T, _ = trans.size()
|
||||
h = 1.0 / data_fps
|
||||
trans_vel = estimate_linear_velocity(trans, h)
|
||||
joints_vel = estimate_linear_velocity(joints3d, h)
|
||||
root_orient_mat = batch_rodrigues(root_orient.reshape((-1, 3))).reshape(
|
||||
(B, T, 3, 3)
|
||||
)
|
||||
root_orient_vel = estimate_angular_velocity(root_orient_mat, h)
|
||||
return trans_vel, joints_vel, root_orient_vel
|
||||
|
||||
|
||||
def estimate_linear_velocity(data_seq, h):
|
||||
"""
|
||||
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) / h
|
||||
init_vel = (data_seq[:, 1:2] - data_seq[:, :1]) / h
|
||||
# middle steps are second order (t+1 - t-1) / 2h
|
||||
middle_vel = (data_seq[:, 2:] - data_seq[:, 0:-2]) / (2 * h)
|
||||
# last step is backward diff (t - t-1) / h
|
||||
final_vel = (data_seq[:, -1:] - data_seq[:, -2:-1]) / h
|
||||
|
||||
vel_seq = torch.cat([init_vel, middle_vel, final_vel], dim=1)
|
||||
return vel_seq
|
||||
|
||||
|
||||
def estimate_angular_velocity(rot_seq, h):
|
||||
"""
|
||||
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, h)
|
||||
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
|
||||
# average 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
|
||||
+643
@@ -0,0 +1,643 @@
|
||||
import os
|
||||
import glob
|
||||
import numpy as np
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import torch.optim as optim
|
||||
|
||||
from body_model import OP_IGNORE_JOINTS
|
||||
from geometry.mesh import save_mesh_scenes, vertices_to_trimesh
|
||||
from util.logger import Logger, log_cur_stats
|
||||
from util.tensor import move_to, detach_all
|
||||
from vis.output import prep_result_vis, animate_scene
|
||||
|
||||
from .losses import RootLoss, SMPLLoss #, MotionLoss
|
||||
from .output import save_camera_json
|
||||
import time
|
||||
|
||||
LINE_SEARCH = "strong_wolfe"
|
||||
LIGHT_BLUE=(0.65098039, 0.74117647, 0.85882353)
|
||||
|
||||
"""
|
||||
Optimization happens in multiple stages
|
||||
Stage 1: fit root orients and trans of each frame indendpently
|
||||
Stage 2: fit SMPL poses and betas of each frame independently
|
||||
Stage 3: fit poses and roots in same global coordinate frame
|
||||
"""
|
||||
|
||||
|
||||
class StageOptimizer(object):
|
||||
def __init__(
|
||||
self,
|
||||
name,
|
||||
model,
|
||||
param_names,
|
||||
lr=1.0,
|
||||
lbfgs_max_iter=20,
|
||||
save_every=10,
|
||||
vis_every=-1,
|
||||
save_meshes=True,
|
||||
max_chunk_steps=10,
|
||||
**kwargs,
|
||||
):
|
||||
Logger.log(f"INITIALIZING OPTIMIZER {name} for {param_names}")
|
||||
self.name = name
|
||||
self.model = model
|
||||
|
||||
self.set_opt_vars(param_names)
|
||||
|
||||
self.optim = torch.optim.LBFGS(
|
||||
self.opt_params, max_iter=lbfgs_max_iter, lr=lr, line_search_fn=LINE_SEARCH
|
||||
)
|
||||
# LBFGS computes losses multiple times per iteration,
|
||||
# save a dict mapping iteration to list of stats_dicts
|
||||
self.loss_dicts = {}
|
||||
|
||||
self.save_every = save_every
|
||||
self.vis_every = vis_every
|
||||
self.save_meshes = save_meshes
|
||||
self.max_chunk_steps = max_chunk_steps
|
||||
|
||||
self.cur_step = 0
|
||||
|
||||
self.add_chunk = 0
|
||||
self.cur_loss = 0
|
||||
self.prev_loss = np.inf
|
||||
self.last_updated = 0
|
||||
self.reached_max = False
|
||||
self.reached_max_iter = -1
|
||||
|
||||
def set_opt_vars(self, param_names):
|
||||
Logger.log("Set param names:")
|
||||
Logger.log(param_names)
|
||||
|
||||
self.param_names = param_names
|
||||
self.model.params.set_require_grads(self.param_names)
|
||||
self.opt_params = [
|
||||
getattr(self.model.params, name) for name in self.param_names
|
||||
]
|
||||
|
||||
def forward_pass(self, obs_data):
|
||||
raise NotImplementedError
|
||||
|
||||
def load_checkpoint(self, out_dir, device=None):
|
||||
if device is None:
|
||||
device = torch.device("cpu")
|
||||
|
||||
param_path = os.path.join(out_dir, f"{self.name}_params.pth")
|
||||
if os.path.isfile(param_path):
|
||||
param_dict = torch.load(param_path, map_location=device)
|
||||
self.model.params.load_dict(param_dict)
|
||||
Logger.log(f"Params loaded from {param_path}")
|
||||
|
||||
optim_path = os.path.join(out_dir, f"{self.name}_optim.pth")
|
||||
if os.path.isfile(optim_path):
|
||||
optim_dict = torch.load(optim_path)
|
||||
self.optim.load_state_dict(optim_dict["optim"])
|
||||
self.cur_step = optim_dict["cur_step"]
|
||||
Logger.log(f"Optimizer loaded from {optim_path} at iter {self.cur_step}")
|
||||
|
||||
def save_checkpoint(self, out_dir):
|
||||
param_path = os.path.join(out_dir, f"{self.name}_params.pth")
|
||||
param_dict = self.model.params.get_dict()
|
||||
if "world_scale" in param_dict:
|
||||
print("WORLD_SCALE", param_dict["world_scale"].detach().cpu())
|
||||
if "floor_plane" in param_dict:
|
||||
print("FLOOR PLANE", param_dict["floor_plane"].detach().cpu())
|
||||
if "cam_f" in param_dict:
|
||||
print("CAM_F", param_dict["cam_f"].detach().cpu())
|
||||
torch.save(param_dict, param_path)
|
||||
Logger.log(f"Model saved at {param_path}")
|
||||
|
||||
optim_path = os.path.join(out_dir, f"{self.name}_optim.pth")
|
||||
torch.save(
|
||||
{"optim": self.optim.state_dict(), "cur_step": self.cur_step},
|
||||
optim_path,
|
||||
)
|
||||
Logger.log(f"Optimizer saved at {optim_path}")
|
||||
|
||||
def save_results(self, out_dir, seq_name):
|
||||
"""
|
||||
pred dict will be a dictionary of trajectories.
|
||||
each trajectory will have params and lists of trimesh sequences
|
||||
"""
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
|
||||
with torch.no_grad():
|
||||
print('go into get_optim_result from optimizer.py save_results')
|
||||
pred_dict = self.model.get_optim_result()
|
||||
pred_dict = move_to(detach_all(pred_dict), "cpu")
|
||||
|
||||
i = self.cur_step
|
||||
for name, results in pred_dict.items():
|
||||
print('save_results, name: ', name)
|
||||
# save parameters of trajectory
|
||||
out_path = f"{out_dir}/{seq_name}_{i:06d}_{name}_results.npz"
|
||||
Logger.log(f"saving params to {out_path}")
|
||||
np.savez(out_path, **results)
|
||||
print('cam_R in save_results and its .npz', results['cam_R'])
|
||||
|
||||
# also save the cameras
|
||||
# print('save the cameras in ptimizer.py save_results')
|
||||
# print(self.model.params)
|
||||
# with torch.no_grad():
|
||||
# cam_R, cam_t = self.model.params.get_extrinsics()
|
||||
# intrins = self.model.params.get_intrinsics()
|
||||
# print('cam_R in 0000.json:', cam_R)
|
||||
# save_camera_json(
|
||||
# f"{out_dir}/{seq_name}_cameras_{self.cur_step:06d}.json",
|
||||
# cam_R.detach().cpu(),
|
||||
# cam_t.detach().cpu(),
|
||||
# intrins.detach().cpu(),
|
||||
# )
|
||||
|
||||
# plot losses
|
||||
self.plot_losses(out_dir)
|
||||
|
||||
# def save_meshes_all(self, res_dir, obs_data, seq_name, num_steps=-1):
|
||||
|
||||
# # check which results are saved
|
||||
# seq_name = obs_data["seq_name"][0]
|
||||
# res_pre = f"{res_dir}/{seq_name}_opt_{self.cur_step:06d}"
|
||||
# with torch.no_grad():
|
||||
# print('go into get_optim_result from optimizer.py vis_result')
|
||||
# pred_dict = self.model.get_optim_result(num_steps=num_steps)
|
||||
|
||||
# res_dict = detach_all(pred_dict["world"])
|
||||
# scene_dict = move_to(
|
||||
# prep_result_vis(
|
||||
# res_dict,
|
||||
# obs_data["vis_mask"],
|
||||
# obs_data["track_id"],
|
||||
# self.model.body_model,
|
||||
# temporal_smooth=True
|
||||
# ),
|
||||
# "cpu",
|
||||
# )
|
||||
|
||||
# # save the meshes
|
||||
# os.makedirs(res_dir, exist_ok=True)
|
||||
# i = self.cur_step
|
||||
# scene_dir = f"{res_dir}/{seq_name}_{i:06d}_meshes"
|
||||
# os.makedirs(scene_dir, exist_ok=True)
|
||||
|
||||
# verts, joints, colors, l_faces, r_faces, is_right, bounds = scene_dict["geometry"]
|
||||
# T = len(verts)
|
||||
# print(f"{T} mesh frames")
|
||||
# times = list(range(0, T, 1))
|
||||
# for t in times:
|
||||
# if len(is_right[t]) > 1:
|
||||
# assert (is_right[t].cpu().numpy().tolist() == [0,1])
|
||||
# # l_meshes = make_batch_mesh(verts[t][0][None], l_faces[t], colors[t][0][None])
|
||||
# # r_meshes = make_batch_mesh(verts[t][1][None], r_faces[t], colors[t][1][None])
|
||||
# # assert len(l_meshes) == 1
|
||||
# # assert len(r_meshes) == 1
|
||||
# tmesh = vertices_to_trimesh(verts[t][0].detach().cpu().numpy(), l_faces[t].detach().cpu().numpy(), LIGHT_BLUE, is_right=0)
|
||||
# tmesh.export(os.path.join(scene_dir, f'{str(t).zfill(6)}_0.obj'))
|
||||
|
||||
# tmesh = vertices_to_trimesh(verts[t][1].detach().cpu().numpy(), r_faces[t].detach().cpu().numpy(), LIGHT_BLUE, is_right=1)
|
||||
# tmesh.export(os.path.join(scene_dir, f'{str(t).zfill(6)}_1.obj'))
|
||||
|
||||
# else:
|
||||
# assert len(is_right[t]) == 1
|
||||
# if is_right[t] == 0:
|
||||
# # meshes = make_batch_mesh(verts[t][0][None], l_faces[t], colors[t][0][None])
|
||||
# tmesh = vertices_to_trimesh(verts[t][0].detach().cpu().numpy(), l_faces[t].detach().cpu().numpy(), LIGHT_BLUE, is_right=0)
|
||||
# tmesh.export(os.path.join(scene_dir, f'{str(t).zfill(6)}_0.obj'))
|
||||
# elif is_right[t] == 1:
|
||||
# # meshes = make_batch_mesh(verts[t][0][None], r_faces[t], colors[t][0][None])
|
||||
# tmesh = vertices_to_trimesh(verts[t][0].detach().cpu().numpy(), r_faces[t].detach().cpu().numpy(), LIGHT_BLUE, is_right=1)
|
||||
# tmesh.export(os.path.join(scene_dir, f'{str(t).zfill(6)}_1.obj'))
|
||||
|
||||
def vis_result(self, res_dir, obs_data, vis=None, num_steps=-1):
|
||||
if vis is None or self.vis_every < 0:
|
||||
return
|
||||
|
||||
# check which results are saved
|
||||
seq_name = obs_data["seq_name"][0]
|
||||
res_pre = f"{res_dir}/{seq_name}_opt_{self.cur_step:06d}"
|
||||
with torch.no_grad():
|
||||
print('go into get_optim_result from optimizer.py vis_result')
|
||||
pred_dict = self.model.get_optim_result(num_steps=num_steps)
|
||||
|
||||
res_dict = detach_all(pred_dict["world"])
|
||||
scene_dict = move_to(
|
||||
prep_result_vis(
|
||||
res_dict,
|
||||
obs_data["vis_mask"],
|
||||
obs_data["track_id"],
|
||||
self.model.body_model,
|
||||
),
|
||||
"cpu",
|
||||
)
|
||||
animate_scene(vis, scene_dict, res_pre, render_views=["src_cam", "above"])
|
||||
|
||||
def log_losses(self, stats_dict):
|
||||
stats_dict = move_to(detach_all(stats_dict), "cpu")
|
||||
log_cur_stats(
|
||||
stats_dict,
|
||||
iter=self.cur_step,
|
||||
to_stdout=(self.cur_step % self.save_every == 0),
|
||||
)
|
||||
for loss_name, loss_val in stats_dict.items():
|
||||
loss_dict = self.loss_dicts.get(loss_name, {})
|
||||
loss_series = loss_dict.get(self.cur_step, [])
|
||||
loss_series.append(loss_val)
|
||||
loss_dict[self.cur_step] = loss_series
|
||||
self.loss_dicts[loss_name] = loss_dict
|
||||
|
||||
def record_current_losses(self, writer):
|
||||
"""
|
||||
record the mean of current step's loss values in tensorboard
|
||||
"""
|
||||
if len(self.loss_dicts) < 1:
|
||||
return
|
||||
|
||||
for loss_name, loss_dict in self.loss_dicts.items():
|
||||
loss_mean = np.mean(loss_dict[self.cur_step])
|
||||
writer.add_scalar(f"{self.name}/{loss_name}", loss_mean, self.cur_step)
|
||||
|
||||
def plot_losses(self, res_dir):
|
||||
"""
|
||||
plot a box plot for each BFGS iteration
|
||||
"""
|
||||
if len(self.loss_dicts) < 1:
|
||||
return
|
||||
for loss_name, loss_dict in self.loss_dicts.items():
|
||||
# times (list len T)
|
||||
# loss vals (list len T of loss value lists)
|
||||
times, loss_vals = zip(*loss_dict.items())
|
||||
plt.figure()
|
||||
plt.boxplot(loss_vals, labels=times, showfliers=False)
|
||||
plt.savefig(f"{res_dir}/{loss_name}.png")
|
||||
|
||||
def run(self, obs_data, num_iters, out_dir, vis=None, writer=None):
|
||||
self.cur_step = 0
|
||||
self.loss.cur_step = 0
|
||||
res_dir = os.path.join(out_dir, self.name)
|
||||
os.makedirs(res_dir, exist_ok=True)
|
||||
seq_name = obs_data["seq_name"][0]
|
||||
print("SEQ NAME", seq_name, "num_iters", num_iters)
|
||||
|
||||
# try to load from checkpoint if exists
|
||||
device = obs_data["joints2d"].device
|
||||
self.load_checkpoint(out_dir, device=device)
|
||||
|
||||
if self.cur_step >= num_iters:
|
||||
Logger.log(f"Current optimizer {self}")
|
||||
Logger.log(f"Checkpoint at {self.cur_step} >= {num_iters}, skipping")
|
||||
return
|
||||
|
||||
Logger.log(f"OPTIMIZING {self.name} FOR {num_iters} ITERATIONS")
|
||||
# time.sleep(5)
|
||||
|
||||
# save initial results and vis
|
||||
print('go into save_results in optimizers.py run()')
|
||||
self.save_results(res_dir, seq_name)
|
||||
|
||||
for i in range(self.cur_step, num_iters):
|
||||
Logger.log("ITER: %d" % (i))
|
||||
|
||||
if (i + 1) % self.save_every == 0: # save before
|
||||
self.save_checkpoint(out_dir)
|
||||
self.save_results(res_dir, seq_name)
|
||||
else:
|
||||
self.save_checkpoint(out_dir)
|
||||
|
||||
if (i + 1) % self.vis_every == 0: # render
|
||||
self.vis_result(res_dir, obs_data, vis)
|
||||
|
||||
self.cur_step = i
|
||||
self.loss.cur_step = i
|
||||
|
||||
#print("ITER: %d" % (i))
|
||||
self.optim_step(obs_data, i, writer)
|
||||
|
||||
# early termination in case of nans
|
||||
if np.isnan(self.cur_loss):
|
||||
# we need to backtrack
|
||||
self.load_checkpoint(out_dir, device=device)
|
||||
print(self.cur_loss, self.kkk, self.ppp)
|
||||
raise ValueError
|
||||
|
||||
# termination for the last chunk
|
||||
if self.reached_max and self.reached_max_iter < 0:
|
||||
self.reached_max_iter = i - 1
|
||||
if self.reached_max and i - self.reached_max_iter >= self.max_chunk_steps:
|
||||
break
|
||||
|
||||
# termination for middle chunks
|
||||
loss_change = self.prev_loss - self.cur_loss
|
||||
if self.last_updated == i - 1 and loss_change == 0:
|
||||
break
|
||||
if (
|
||||
(self.cur_loss < 0 and loss_change < 100)
|
||||
or (i - self.last_updated >= self.max_chunk_steps)
|
||||
or (loss_change < 20 and i - self.last_updated > 5)
|
||||
):
|
||||
self.add_chunk = self.add_chunk + 1
|
||||
self.last_updated = i
|
||||
self.prev_loss = self.cur_loss
|
||||
|
||||
# final save and vis step
|
||||
self.cur_step = num_iters
|
||||
self.save_checkpoint(out_dir)
|
||||
self.save_results(res_dir, seq_name)
|
||||
self.vis_result(res_dir, obs_data, vis)
|
||||
# self.save_meshes_all(res_dir, obs_data, seq_name)
|
||||
|
||||
def optim_step(self, obs_data, i, writer=None):
|
||||
def closure():
|
||||
self.optim.zero_grad()
|
||||
#print('******************************', i)
|
||||
loss, stats_dict, preds = self.forward_pass(obs_data)
|
||||
#print('******************************', i)
|
||||
stats_dict["total"] = loss
|
||||
self.log_losses(move_to(detach_all(stats_dict), "cpu"))
|
||||
self.cur_loss = stats_dict["total"].detach().cpu().item()
|
||||
self.kkk = stats_dict
|
||||
self.ppp = preds
|
||||
loss_keys = stats_dict.keys()
|
||||
if "motion_prior" in loss_keys and "pose_prior" in loss_keys:
|
||||
self.cur_loss = (
|
||||
self.cur_loss
|
||||
- 0.04 * stats_dict["pose_prior"].detach().cpu().item()
|
||||
)
|
||||
print("optimizer print", loss, loss.shape, stats_dict, loss.requires_grad)
|
||||
loss.backward()
|
||||
# print('end of loss backward')
|
||||
return loss
|
||||
|
||||
#print('!!!!!!!!!!!!!!!!!!!!!!!!!!', i)
|
||||
self.optim.step(closure)
|
||||
#print('!!!!!!!!!!!!!!!!!!!!!!!!!!', i)
|
||||
if writer is not None:
|
||||
self.record_current_losses(writer)
|
||||
|
||||
|
||||
class RootOptimizer(StageOptimizer):
|
||||
name = "root_fit"
|
||||
stage = 0
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model,
|
||||
all_loss_weights,
|
||||
use_chamfer=False,
|
||||
robust_loss_type="none",
|
||||
robust_tuning_const=4.6851,
|
||||
joints2d_sigma=100,
|
||||
**kwargs,
|
||||
):
|
||||
param_names = ["trans", "root_orient"]
|
||||
super().__init__(self.name, model, param_names, **kwargs)
|
||||
|
||||
self.loss = RootLoss(
|
||||
all_loss_weights[self.stage],
|
||||
ignore_op_joints=OP_IGNORE_JOINTS,
|
||||
joints2d_sigma=joints2d_sigma,
|
||||
use_chamfer=use_chamfer,
|
||||
robust_loss=robust_loss_type,
|
||||
robust_tuning_const=robust_tuning_const,
|
||||
faces=model.body_model.faces_tensor
|
||||
)
|
||||
|
||||
def forward_pass(self, obs_data):
|
||||
"""
|
||||
Takes in observed data, predicts the smpl parameters and returns loss
|
||||
"""
|
||||
# raise ValueError
|
||||
# Use current params to go through SMPL and get joints3d, verts3d, points3d
|
||||
pred_data = self.model.pred_params_mano(obs_data["is_right"])
|
||||
pred_data["cameras"] = self.model.params.get_cameras()
|
||||
|
||||
# compute data losses only
|
||||
vis_mask = obs_data["vis_mask"] >= 0
|
||||
# print('entering loss from RootOptimizer')
|
||||
loss, stats_dict = self.loss(obs_data, pred_data, vis_mask)
|
||||
return loss, stats_dict, pred_data
|
||||
|
||||
|
||||
class SmoothOptimizer(StageOptimizer):
|
||||
name = "smooth_fit"
|
||||
stage = 1
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model,
|
||||
all_loss_weights,
|
||||
use_chamfer=False,
|
||||
robust_loss_type="none",
|
||||
robust_tuning_const=4.6851,
|
||||
joints2d_sigma=100,
|
||||
**kwargs,
|
||||
):
|
||||
param_names = ["trans", "root_orient", "betas", "latent_pose"]
|
||||
if model.opt_scale:
|
||||
param_names += ["world_scale"]
|
||||
if model.opt_cams:
|
||||
param_names += ["cam_f", "delta_cam_R"]
|
||||
|
||||
super().__init__(self.name, model, param_names, **kwargs)
|
||||
|
||||
self.loss = SMPLLoss(
|
||||
all_loss_weights[self.stage],
|
||||
ignore_op_joints=OP_IGNORE_JOINTS,
|
||||
joints2d_sigma=joints2d_sigma,
|
||||
use_chamfer=use_chamfer,
|
||||
robust_loss=robust_loss_type,
|
||||
robust_tuning_const=robust_tuning_const,
|
||||
faces=model.body_model.faces_tensor
|
||||
)
|
||||
|
||||
def forward_pass(self, obs_data):
|
||||
"""
|
||||
Takes in observed data, predicts the smpl parameters and returns loss
|
||||
"""
|
||||
# raise ValueError
|
||||
# Use current params to go through SMPL and get joints3d, verts3d, points3d
|
||||
pred_data = self.model.pred_params_mano(obs_data["is_right"])
|
||||
pred_data["cameras"] = self.model.params.get_cameras()
|
||||
pred_data.update(self.model.params.get_vars())
|
||||
|
||||
# camera predictions
|
||||
# pts_ij, target, weight = self.model.params.get_reprojected_points()
|
||||
# pred_data["bg2d_err"] = (
|
||||
# (weight * (pts_ij - target) ** 2).sum(dim=(-1, -2)).mean()
|
||||
# )
|
||||
pred_data["cam_R"], pred_data["cam_t"] = self.model.params.get_extrinsics()
|
||||
|
||||
# compute data losses only
|
||||
vis_mask = obs_data["vis_mask"] >= 0
|
||||
loss, stats_dict = self.loss(obs_data, pred_data, self.model.seq_len, vis_mask)
|
||||
return loss, stats_dict, pred_data
|
||||
|
||||
|
||||
# class MotionOptimizer(StageOptimizer):
|
||||
# name = "motion_fit"
|
||||
# stage = 2
|
||||
|
||||
# def __init__(self, model, all_loss_weights, opt_cams=False, **kwargs):
|
||||
# self.opt_cams = opt_cams and model.opt_cams
|
||||
# param_names = [
|
||||
# "trans",
|
||||
# "root_orient",
|
||||
# "latent_pose",
|
||||
# "trans_vel",
|
||||
# "root_orient_vel",
|
||||
# "joints_vel",
|
||||
# "betas",
|
||||
# "latent_motion",
|
||||
# "floor_plane",
|
||||
# ]
|
||||
# if model.opt_scale:
|
||||
# param_names += ["world_scale"]
|
||||
|
||||
# if self.opt_cams:
|
||||
# Logger.log(f"{self.name} OPTIMIZING CAMERAS")
|
||||
# param_names += ["delta_cam_R", "cam_f"]
|
||||
|
||||
# super().__init__(self.name, model, param_names, **kwargs)
|
||||
# self.set_loss(model, all_loss_weights[self.stage], **kwargs)
|
||||
|
||||
# def set_loss(
|
||||
# self,
|
||||
# model,
|
||||
# loss_weights,
|
||||
# use_chamfer=False,
|
||||
# robust_loss_type="none",
|
||||
# robust_tuning_const=4.6851,
|
||||
# joints2d_sigma=100,
|
||||
# **kwargs,
|
||||
# ):
|
||||
# self.loss = MotionLoss(
|
||||
# loss_weights,
|
||||
# init_motion_prior=model.init_motion_prior,
|
||||
# ignore_op_joints=OP_IGNORE_JOINTS,
|
||||
# joints2d_sigma=joints2d_sigma,
|
||||
# use_chamfer=use_chamfer,
|
||||
# robust_loss=robust_loss_type,
|
||||
# robust_tuning_const=robust_tuning_const,
|
||||
# )
|
||||
|
||||
# def get_motion_scale(self):
|
||||
# return 1.0
|
||||
|
||||
# def forward_pass(self, obs_data, num_steps=-1):
|
||||
# p = self.model.params
|
||||
# param_names = [
|
||||
# "betas",
|
||||
# "joints_vel",
|
||||
# "trans_vel",
|
||||
# "root_orient_vel",
|
||||
# ]
|
||||
# param_dict = p.get_vars(param_names)
|
||||
# # param_dict["floor_plane"] = p.floor_plane[p.floor_idcs]
|
||||
|
||||
# preds, world_preds = self.model.rollout_mano_steps(num_steps)
|
||||
# preds.update(param_dict)
|
||||
# world_preds.update(param_dict)
|
||||
|
||||
# # track mask is the length of the sequence that contains num_steps of each track
|
||||
# track_mask = world_preds.get("track_mask", None)
|
||||
# T = track_mask.shape[1] if track_mask is not None else num_steps
|
||||
# world_preds["cameras"] = p.get_cameras(np.arange(T))
|
||||
# if self.opt_cams:
|
||||
# cam_R, cam_t = p.get_extrinsics()
|
||||
# world_preds["cam_R"], world_preds["cam_t"] = cam_R[:T], cam_t[:T]
|
||||
|
||||
# obs_data = slice_dict(obs_data, 0, T)
|
||||
# motion_scale = self.get_motion_scale()
|
||||
# loss, stats_dict = self.loss(
|
||||
# obs_data,
|
||||
# preds,
|
||||
# world_preds,
|
||||
# self.model.seq_len,
|
||||
# valid_mask=track_mask,
|
||||
# init_motion_scale=motion_scale,
|
||||
# )
|
||||
# return loss, stats_dict, (preds, world_preds)
|
||||
|
||||
|
||||
# def slice_dict(d, start, end):
|
||||
# out = d.copy()
|
||||
# for k, v in d.items():
|
||||
# if not isinstance(v, torch.Tensor) or v.ndim < 3 or v.shape[1] < end:
|
||||
# continue
|
||||
# out[k] = v[:, start:end]
|
||||
# return out
|
||||
|
||||
|
||||
# class MotionOptimizerChunks(MotionOptimizer):
|
||||
# """
|
||||
# Incrementally optimize sequence in chunks (with all variables)
|
||||
# :param chunk_size (int) number of frames per chunk
|
||||
# :param init_steps (int) number of optimization steps to spend on first chunk
|
||||
# :param chunk_steps (int) number of opt steps to spend on subsequent chunks
|
||||
# """
|
||||
|
||||
# name = "motion_chunks"
|
||||
# stage = 2
|
||||
|
||||
# def __init__(self, *args, chunk_size=20, init_steps=20, chunk_steps=10, **kwargs):
|
||||
# self.chunk_size = chunk_size
|
||||
# self.init_steps = init_steps
|
||||
# self.chunk_steps = chunk_steps
|
||||
# super().__init__(*args, **kwargs)
|
||||
|
||||
# @property
|
||||
# def num_iters(self):
|
||||
# num_chunks = int(np.ceil(self.model.seq_len / self.chunk_size)) + 1
|
||||
# return self.init_steps + self.chunk_steps * num_chunks
|
||||
|
||||
# @property
|
||||
# def end_idx(self):
|
||||
# if self.cur_step < self.init_steps and self.add_chunk == 0:
|
||||
# return self.chunk_size
|
||||
# chunk_idx = max(
|
||||
# (self.cur_step - self.init_steps) // self.chunk_steps + 1, self.add_chunk
|
||||
# )
|
||||
# if int(np.ceil(self.model.seq_len / self.chunk_size)) == self.add_chunk:
|
||||
# self.reached_max = True
|
||||
# return min(self.chunk_size * (chunk_idx + 1), self.model.seq_len)
|
||||
|
||||
# @property
|
||||
# def start_idx(self):
|
||||
# return 0
|
||||
|
||||
# def get_motion_scale(self):
|
||||
# num_frames = self.end_idx - self.start_idx
|
||||
# return max(1.0, float(self.model.seq_len) / num_frames)
|
||||
|
||||
# def forward_pass(self, obs_data):
|
||||
# return super().forward_pass(obs_data, num_steps=self.end_idx)
|
||||
|
||||
# def vis_result(self, res_dir, obs_data, vis=None, **kwargs):
|
||||
# print("start, end", self.start_idx, self.end_idx)
|
||||
# return super().vis_result(res_dir, obs_data, vis=vis, num_steps=self.end_idx)
|
||||
|
||||
|
||||
# class MotionOptimizerFreeze(MotionOptimizer):
|
||||
# """
|
||||
# Only optimizing a subset of parameters
|
||||
# """
|
||||
|
||||
# name = "motion_freeze"
|
||||
# stage = 2
|
||||
|
||||
# def __init__(self, model, all_loss_weights, no_contacts=True, **kwargs):
|
||||
# loss_weights = all_loss_weights[self.stage].copy()
|
||||
|
||||
# if no_contacts:
|
||||
# loss_weights["contact_height"] = 0.0
|
||||
# loss_weights["contact_vel"] = 0.0
|
||||
|
||||
# param_names = ["latent_motion", "betas", "trans", "floor_plane"]
|
||||
# if model.optim_scale:
|
||||
# param_names += ["world_scale"]
|
||||
|
||||
# StageOptimizer.__init__(self, self.name, model, param_names, **kwargs)
|
||||
# self.set_loss(model, loss_weights, **kwargs)
|
||||
+188
@@ -0,0 +1,188 @@
|
||||
import os
|
||||
import imageio
|
||||
import glob
|
||||
import json
|
||||
import subprocess
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from util.tensor import move_to, detach_all, to_torch
|
||||
from util.logger import Logger
|
||||
|
||||
|
||||
def get_results_paths(res_dir):
|
||||
"""
|
||||
get the iterations of all saved results in res_dir
|
||||
:param res_dir (str) result dir
|
||||
returns a dict of iter to result path
|
||||
"""
|
||||
res_files = sorted(glob.glob(f"{res_dir}/*_results.npz"))
|
||||
print(f"found {len(res_files)} results in {res_dir}")
|
||||
|
||||
path_dict = {}
|
||||
for res_file in res_files:
|
||||
it, name, _ = os.path.basename(res_file).split("_")[-3:]
|
||||
assert name in ["world", "prior"]
|
||||
if it not in path_dict:
|
||||
path_dict[it] = {}
|
||||
path_dict[it][name] = res_file
|
||||
return path_dict
|
||||
|
||||
|
||||
def load_result(res_path_dict):
|
||||
"""
|
||||
load all saved results for a given iteration
|
||||
:param res_path_dict (dict) paths to relevant results
|
||||
returns dict of results
|
||||
"""
|
||||
res_dict = {}
|
||||
for name, path in res_path_dict.items():
|
||||
res = np.load(path)
|
||||
res_dict[name] = to_torch({k: res[k] for k in res.files})
|
||||
return res_dict
|
||||
|
||||
|
||||
def save_initial_predictions(model, out_dir, seq_name):
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
with torch.no_grad():
|
||||
print('go into get_optim_result from output.py save_initial_predictions')
|
||||
pred_dict = model.get_optim_result()
|
||||
pred_dict = move_to(detach_all(pred_dict), "cpu")
|
||||
|
||||
for name, results in pred_dict.items():
|
||||
out_path = f"{out_dir}/{seq_name}_{0:06d}_init_{name}_results.npz"
|
||||
Logger.log(f"saving initial results to {out_path}")
|
||||
np.savez(out_path, **results)
|
||||
# print('cam_R in save_initial_predictions: ', results['cam_R'].shape, results['cam_R'])
|
||||
# print((results['cam_R'][0] == results['cam_R'][1]).all())
|
||||
# print((results['cam_R'][1] == results['cam_R'][2]).all())
|
||||
# print((results['cam_R'][2] == results['cam_R'][3]).all())
|
||||
# print((results['cam_R'][3] == results['cam_R'][4]).all())
|
||||
# print((results['cam_R'][4] == results['cam_R'][5]).all())
|
||||
# print((results['cam_R'][5] == results['cam_R'][6]).all())
|
||||
|
||||
|
||||
def save_input_poses(dataset, out_dir, seq_name, name="phalp"):
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
dataset.load_data(interp_input=False)
|
||||
d = dataset.data_dict
|
||||
res = {
|
||||
"pose_body": np.stack(d["init_body_pose"], axis=0),
|
||||
"trans": np.stack(d["init_trans"], axis=0),
|
||||
"root_orient": np.stack(d["init_root_orient"], axis=0),
|
||||
}
|
||||
print({k: v.shape for k, v in res.items()})
|
||||
out_path = f"{out_dir}/{seq_name}_{0:06d}_{name}_world_results.npz"
|
||||
Logger.log(f"saving inputs to {out_path}")
|
||||
np.savez(out_path, **res)
|
||||
|
||||
|
||||
def save_input_frames_ffmpeg(dataset, out_dir, name="input", fps=30, overwrite=False):
|
||||
vid_path = f"{out_dir}/{name}.mp4"
|
||||
if not overwrite and os.path.isfile(vid_path):
|
||||
return
|
||||
|
||||
list_path = f"{out_dir}/src_paths.txt"
|
||||
with open(list_path, "w") as f:
|
||||
f.write("\n".join([f"file '{p}'" for p in dataset.sel_img_paths]))
|
||||
|
||||
filter_str = f"-framerate {fps} -vf 'format=yuv420p'"
|
||||
cmd = (
|
||||
f"ffmpeg -loglevel error -f concat -safe 0 "
|
||||
f"-i {list_path} -c:v libx264 {filter_str} {vid_path} -y"
|
||||
)
|
||||
print(cmd)
|
||||
subprocess.call(cmd, shell=True, stdin=subprocess.PIPE)
|
||||
print(f"SAVED {len(dataset.sel_img_paths)} INPUT FRAMES TO {vid_path}")
|
||||
return vid_path
|
||||
|
||||
|
||||
def save_input_frames(dataset, vid_path, fps=30, overwrite=False):
|
||||
if not overwrite and os.path.isfile(vid_path):
|
||||
return
|
||||
|
||||
writer = imageio.get_writer(vid_path, fps=fps)
|
||||
for path in dataset.sel_img_paths:
|
||||
writer.append_data(imageio.imread(path))
|
||||
writer.close()
|
||||
print(f"SAVED {len(dataset.sel_img_paths)} INPUT FRAMES TO {vid_path}")
|
||||
return vid_path
|
||||
|
||||
|
||||
def load_track_info(path):
|
||||
assert os.path.isfile(path)
|
||||
|
||||
with open(path, "r") as f:
|
||||
track_info = json.load(f)
|
||||
|
||||
meta = track_info["meta"]
|
||||
tracks = track_info["tracks"]
|
||||
track_ids, idcs, vis_masks = map(
|
||||
torch.as_tensor,
|
||||
zip(
|
||||
*[
|
||||
(int(tid), info["index"], info["vis_mask"])
|
||||
for tid, info in tracks.items()
|
||||
]
|
||||
),
|
||||
)
|
||||
track_ids = track_ids[idcs]
|
||||
vis_masks = vis_masks[idcs]
|
||||
data_interval = meta["data_interval"]
|
||||
seq_interval = meta["seq_interval"]
|
||||
return track_ids, vis_masks, data_interval, seq_interval
|
||||
|
||||
|
||||
def save_track_info(dataset, out_dir):
|
||||
# track indices are relative to the entire sequence
|
||||
track_ids = dataset.track_ids
|
||||
track_dict = {}
|
||||
for i, (tid, mask) in enumerate(zip(dataset.track_ids, dataset.track_vis_masks)):
|
||||
track_dict[int(tid)] = {"index": i, "vis_mask": mask.tolist()}
|
||||
|
||||
out_dict = {
|
||||
"tracks": track_dict,
|
||||
"meta": {
|
||||
"seq_interval": (int(dataset.start_idx), int(dataset.end_idx)),
|
||||
"data_interval": (int(dataset.data_start), int(dataset.data_end)),
|
||||
},
|
||||
}
|
||||
|
||||
with open(f"{out_dir}/track_info.json", "w") as f:
|
||||
json.dump(out_dict, f)
|
||||
print("SAVED TRACK INFO")
|
||||
|
||||
|
||||
def load_camera_json(path):
|
||||
assert os.path.isfile(path)
|
||||
|
||||
with open(path, "r") as f:
|
||||
cam_data = json.load(f)
|
||||
R = torch.as_tensor(cam_data["rotation"]).reshape(-1, 3, 3)
|
||||
t = torch.as_tensor(cam_data["translation"])
|
||||
intrins = torch.as_tensor(cam_data["intrinsics"])
|
||||
return R.float(), t.float(), intrins.float()
|
||||
|
||||
|
||||
def save_camera_json(path, cam_R, cam_t, intrins):
|
||||
"""
|
||||
:param path
|
||||
:param cam_R (N, 3, 3)
|
||||
:param cam_t (N, 3)
|
||||
:param intrins (N, 4)
|
||||
"""
|
||||
N = len(cam_R)
|
||||
T = torch.tensor([[1, 0, 0], [0, -1, 0], [0, 0, -1]], dtype=torch.float32)
|
||||
cam_R = torch.einsum("ij,bjk->bik", T, cam_R)
|
||||
cam_t = torch.einsum("ij,bj->bi", T, cam_t)
|
||||
with open(path, "w") as f:
|
||||
json.dump(
|
||||
{
|
||||
"rotation": cam_R.reshape(N, 9).tolist(),
|
||||
"translation": cam_t.tolist(),
|
||||
"intrinsics": intrins.tolist(),
|
||||
},
|
||||
f,
|
||||
indent=1,
|
||||
)
|
||||
+174
@@ -0,0 +1,174 @@
|
||||
import numpy as np
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
# from body_model import MANO_JOINTS
|
||||
from geometry.camera import invert_camera, compose_cameras
|
||||
from geometry.rotation import batch_rodrigues
|
||||
from util.logger import Logger
|
||||
from util.tensor import detach_all
|
||||
|
||||
# J_BODY = len(MANO_JOINTS) - 1 # no root
|
||||
# assert J_BODY == 15
|
||||
|
||||
class Params(nn.Module):
|
||||
def __init__(self, batch_size):
|
||||
super().__init__()
|
||||
self.batch_size = batch_size
|
||||
self.param_names = set()
|
||||
|
||||
def set_param(self, name, val, requires_grad=False):
|
||||
print("SETTING PARAM", name, val.shape)
|
||||
with torch.no_grad():
|
||||
setattr(
|
||||
self, name, nn.Parameter(val.contiguous(), requires_grad=requires_grad)
|
||||
)
|
||||
self.param_names.add(name)
|
||||
|
||||
def get_param(self, name):
|
||||
if name not in self.param_names:
|
||||
raise ValueError(f"{name} not stored as opt param")
|
||||
return getattr(self, name)
|
||||
|
||||
def load_dict(self, param_dict):
|
||||
for name, val in param_dict.items():
|
||||
self.set_param(name, val, requires_grad=False)
|
||||
|
||||
def get_dict(self):
|
||||
return {name: self.get_param_item(name) for name in self.param_names}
|
||||
|
||||
def get_vars(self, names=None):
|
||||
if names is None:
|
||||
names = self.param_names
|
||||
return {name: self.get_param(name) for name in names}
|
||||
|
||||
def get_param_item(self, name):
|
||||
with torch.no_grad():
|
||||
param = self.get_param(name)
|
||||
return param.detach()
|
||||
|
||||
def _set_param_grad(self, name, val: bool):
|
||||
if name not in self.param_names:
|
||||
raise ValueError(f"{name} not stored as param")
|
||||
param = getattr(self, name)
|
||||
assert isinstance(param, torch.Tensor)
|
||||
param.requires_grad = val
|
||||
|
||||
def set_require_grads(self, names):
|
||||
"""
|
||||
set parameters in names to True, set all others to False
|
||||
"""
|
||||
for name in self.param_names:
|
||||
self._set_param_grad(name, False)
|
||||
|
||||
for name in names:
|
||||
self._set_param_grad(name, True)
|
||||
|
||||
Logger.log("Set parameter grads:")
|
||||
Logger.log(
|
||||
{name: getattr(self, name).requires_grad for name in self.param_names}
|
||||
)
|
||||
|
||||
|
||||
class CameraParams(Params):
|
||||
"""
|
||||
Parameter container with cameras
|
||||
"""
|
||||
|
||||
def set_cameras(
|
||||
self, cam_data, opt_scale=True, opt_cams=False, opt_focal=False, **kwargs
|
||||
):
|
||||
self.opt_scale = opt_scale
|
||||
self.opt_cams = opt_cams
|
||||
self.opt_focal = opt_focal
|
||||
|
||||
# (T, 3, 3), (T, 3)
|
||||
cam_R, cam_t = cam_data["cam_R"], cam_data["cam_t"]
|
||||
intrins = cam_data["intrins"] # (T, 4)
|
||||
|
||||
T = cam_R.shape[0]
|
||||
device = cam_R.device
|
||||
|
||||
# assume focal length and center are same for all timesteps
|
||||
# print(intrins.shape, cam_R.shape, cam_t.shape)
|
||||
# torch.Size([128, 4]) torch.Size([128, 3, 3]) torch.Size([128, 3])
|
||||
|
||||
self.cam_center = intrins[:, 2:] # (T, 2)
|
||||
cam_f = intrins[:, :2] # (T, 2)
|
||||
if self.opt_focal:
|
||||
self.set_param("cam_f", cam_f)
|
||||
else:
|
||||
self.cam_f = cam_f
|
||||
|
||||
self._cam_R = cam_R # (T, 3, 3)
|
||||
self._cam_t = cam_t # (T, 3)
|
||||
|
||||
world_scale = torch.ones(1, 1, device=device)
|
||||
if self.opt_scale:
|
||||
if "world_scale" in kwargs:
|
||||
world_scale = kwargs["world_scale"]
|
||||
self.set_param("world_scale", world_scale)
|
||||
else:
|
||||
self.world_scale = world_scale
|
||||
|
||||
if self.opt_cams:
|
||||
init_delta_R = torch.zeros(T, 3, device=device)
|
||||
init_delta_t = torch.zeros(T, 3, device=device)
|
||||
if "delta_cam_R" in kwargs:
|
||||
print("setting delta_cam_R from kwargs")
|
||||
init_delta_R = kwargs["delta_cam_R"]
|
||||
if "delta_cam_t" in kwargs:
|
||||
print("setting delta_cam_t from kwargs")
|
||||
init_delta_t = kwargs["delta_cam_t"]
|
||||
|
||||
self.set_param("delta_cam_R", init_delta_R)
|
||||
self.set_param("delta_cam_t", init_delta_t)
|
||||
|
||||
@property
|
||||
def intrins(self): # (4,)
|
||||
return torch.cat([self.cam_f[0], self.cam_center[0]], dim=-1).detach().cpu()
|
||||
|
||||
def get_extrinsics(self):
|
||||
"""
|
||||
returns (T, 3, 3), (T, 3)
|
||||
"""
|
||||
cam_R, cam_t = self._cam_R, self._cam_t
|
||||
print('get_extrinsics')
|
||||
if self.opt_cams:
|
||||
raise ValueError
|
||||
print('self.opt_cams:', self.opt_cams)
|
||||
# raise ValueError
|
||||
dR = batch_rodrigues(self.delta_cam_R)
|
||||
cam_R = torch.matmul(cam_R, dR)
|
||||
cam_t = cam_t + self.delta_cam_t
|
||||
return cam_R, cam_t * self.world_scale
|
||||
|
||||
def get_intrinsics(self):
|
||||
"""
|
||||
returns (T, 4)
|
||||
"""
|
||||
return torch.cat([self.cam_f, self.cam_center], axis=-1)
|
||||
|
||||
def get_cameras(self, idcs=None):
|
||||
"""
|
||||
returns cam_R (B, T, 3, 3) cam_t (B, T, 3), cam_f (T, 2), cam_center (T, 2)
|
||||
"""
|
||||
if idcs is None:
|
||||
idcs = np.arange(self._cam_R.shape[0])
|
||||
|
||||
# print('idcs:', idcs, self._cam_R.shape[0], self._cam_R.shape)
|
||||
cam_R, cam_t = self.get_extrinsics()
|
||||
cam_R, cam_t = cam_R[None, idcs], cam_t[None, idcs]
|
||||
cam_f, cam_center = self.cam_f[idcs], self.cam_center[idcs]
|
||||
|
||||
B = self.batch_size
|
||||
cam_R = cam_R.repeat(B, 1, 1, 1)
|
||||
cam_t = cam_t.repeat(B, 1, 1)
|
||||
# print(cam_R.shape, cam_t.shape, cam_f.shape, cam_center.shape)
|
||||
# print(B, self.batch_size)
|
||||
# raise ValueError
|
||||
# torch.Size([2, 128, 3, 3]) torch.Size([2, 128, 3]) torch.Size([128, 2]) torch.Size([128, 2])
|
||||
# 2 2
|
||||
return cam_R, cam_t, cam_f, cam_center
|
||||
+137
@@ -0,0 +1,137 @@
|
||||
import os
|
||||
import glob
|
||||
import pandas as pd
|
||||
from copy import deepcopy
|
||||
import warnings
|
||||
|
||||
|
||||
def update_args(args):
|
||||
args = deepcopy(args)
|
||||
if args.type == "egobody":
|
||||
if args.root is None:
|
||||
warnings.warn(
|
||||
"args.root flag is set to None, using root directory in source"
|
||||
)
|
||||
args.root = "/path/to/egobody"
|
||||
if args.seqs is None:
|
||||
args.seqs = get_egobody_seqs(args.root, args.split)
|
||||
args.img_dirs = [
|
||||
glob.glob(f"{args.root}/egocentric_color/{seq}/**/PV")[0]
|
||||
for seq in args.seqs
|
||||
]
|
||||
return args
|
||||
|
||||
if args.type == "3dpw":
|
||||
if args.root is None:
|
||||
warnings.warn(
|
||||
"args.root flag is set to None, using root directory in source"
|
||||
)
|
||||
args.root = "/path/to/3DPW"
|
||||
if args.split == "val":
|
||||
args.split = "validation"
|
||||
if args.seqs is None:
|
||||
args.seqs = get_3dpw_seqs(args.root, args.split)
|
||||
args.img_name = "imageFiles"
|
||||
|
||||
elif args.type == "posetrack":
|
||||
if args.root is None:
|
||||
warnings.warn(
|
||||
"args.root flag is set to None, using root directory in source"
|
||||
)
|
||||
args.root = "/path/to/posetrack/posetrack2018/posetrack_data"
|
||||
args.img_name = f"images/{args.split}"
|
||||
|
||||
elif args.type == "davis":
|
||||
if args.root is None:
|
||||
warnings.warn(
|
||||
"args.root flag is set to None, using root directory in source"
|
||||
)
|
||||
args.root = "/path/to/DAVIS"
|
||||
args.img_name = "JPEGImages/Full-Resolution"
|
||||
args.split = ""
|
||||
|
||||
else:
|
||||
assert args.root is not None
|
||||
args.img_name = args.img_name if args.img_name is not None else "images"
|
||||
args.split = ""
|
||||
|
||||
args.img_dirs, args.seqs = get_all_img_dirs(
|
||||
f"{args.root}/{args.img_name}", args.seqs
|
||||
)
|
||||
return args
|
||||
|
||||
|
||||
def isimage(f):
|
||||
ext = os.path.splitext(f)[-1].lower()
|
||||
return ext in [".png", ".jpg"]
|
||||
|
||||
|
||||
def get_all_img_dirs(data_root, seqs=None):
|
||||
"""
|
||||
returns all image directories in root, recursively searching, and the sequence names
|
||||
"""
|
||||
seqs = seqs if seqs is not None else []
|
||||
img_dirs = []
|
||||
|
||||
def get_name(p):
|
||||
return p.removeprefix(f"{data_root}/")
|
||||
|
||||
for root, _, files in os.walk(data_root):
|
||||
num_imgs = len([f for f in files if isimage(f)])
|
||||
if num_imgs > 0:
|
||||
name = get_name(root)
|
||||
if len(seqs) == 0 or name in seqs:
|
||||
img_dirs.append(root)
|
||||
img_dirs = sorted(img_dirs)
|
||||
seqs = [get_name(d) for d in img_dirs]
|
||||
return img_dirs, seqs
|
||||
|
||||
|
||||
def get_custom_seqs(data_root):
|
||||
img_dir = f"{data_root}/images"
|
||||
if not os.path.isdir(img_dir):
|
||||
return []
|
||||
return sorted(os.listdir(img_dir))
|
||||
|
||||
|
||||
def get_egobody_seqs(data_root, split):
|
||||
split_file = f"{data_root}/data_splits.csv"
|
||||
df = pd.read_csv(split_file)
|
||||
if split not in df.columns:
|
||||
print(f"{split} not in {split_file}")
|
||||
return []
|
||||
return sorted(df[split].dropna().tolist())
|
||||
|
||||
|
||||
def get_3dpw_seqs(data_root, split):
|
||||
split_dir = f"{data_root}/sequenceFiles/{split}"
|
||||
if not os.path.isdir(split_dir):
|
||||
return []
|
||||
seq_files = sorted(os.listdir(split_dir))
|
||||
return [os.path.splitext(f)[0] for f in seq_files]
|
||||
|
||||
|
||||
def get_posetrack_seqs(data_root, split):
|
||||
split_dir = f"{data_root}/images/{split}"
|
||||
if not os.path.isdir(split_dir):
|
||||
return []
|
||||
return sorted(os.listdir(split_dir))
|
||||
|
||||
|
||||
def get_davis_seqs(data_root):
|
||||
img_root = f"{data_root}/JPEGImages/Full-Resolution"
|
||||
if not os.path.isdir(img_root):
|
||||
return []
|
||||
return sorted(os.listdir(img_root))
|
||||
|
||||
|
||||
def get_img_dir(data_type, data_root, seq, img_name):
|
||||
if data_type == "posetrack":
|
||||
return f"{data_root}/{img_name}/{seq}"
|
||||
if data_type == "3dpw":
|
||||
return f"{data_root}/imageFiles/{seq}"
|
||||
if data_type == "davis":
|
||||
return f"{data_root}/JPEGImages/Full-Resolution/{seq}"
|
||||
if data_type == "custom":
|
||||
return f"{data_root}/images/{seq}" # custom sequence
|
||||
return None
|
||||
@@ -0,0 +1,124 @@
|
||||
import os
|
||||
import joblib
|
||||
import json
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
|
||||
def export_hamer_predictions(res_path, target_dir):
|
||||
"""
|
||||
exports hamer output results into a directory for each track
|
||||
of JSON files with each frame's mano prediction
|
||||
"""
|
||||
os.makedirs(target_dir, exist_ok=True)
|
||||
print(f"exporting hamer predictions from {res_path} to {target_dir}")
|
||||
|
||||
tracklet_data = joblib.load(res_path)
|
||||
for frame, track_data in tracklet_data.items():
|
||||
name = os.path.splitext(frame.split("/")[-1])[0]
|
||||
print(track_data['tid'])
|
||||
track_dicts = unpack_frame(track_data)
|
||||
for tid, pred_dict in track_dicts.items():
|
||||
track_dir = os.path.join(target_dir, f"{tid:03d}")
|
||||
os.makedirs(track_dir, exist_ok=True)
|
||||
pred_path = os.path.join(track_dir, f"{name}_mano.json")
|
||||
with open(pred_path, "w") as f:
|
||||
json.dump(pred_dict, f, indent=1)
|
||||
|
||||
|
||||
def export_vitpose_keypoints(res_path, target_dir):
|
||||
"""
|
||||
exports each track of each frame into its own json of keypoints
|
||||
"""
|
||||
os.makedirs(target_dir, exist_ok=True)
|
||||
print(f"exporting keypoints from {res_path} to {target_dir}")
|
||||
|
||||
tracklet_data = joblib.load(res_path)
|
||||
for frame, track_data in tracklet_data.items():
|
||||
name = os.path.splitext(frame.split("/")[-1])[0]
|
||||
tids = track_data["tid"]
|
||||
valid_tids = track_data["tracked_ids"]
|
||||
kps_all = track_data["extra_data"]
|
||||
# mask_paths = track_data["mask_name"]
|
||||
for i, tid in enumerate(tids):
|
||||
if tid not in valid_tids:
|
||||
continue
|
||||
track_dir = os.path.join(target_dir, f"{tid:03d}")
|
||||
os.makedirs(track_dir, exist_ok=True)
|
||||
kp_path = os.path.join(track_dir, f"{name}_keypoints.json")
|
||||
kp_dict = {
|
||||
"people": [
|
||||
{
|
||||
"pose_keypoints_2d": kps_all[i],#.tolist(),
|
||||
# "mask_path": mask_paths[i],
|
||||
}
|
||||
],
|
||||
}
|
||||
with open(kp_path, "w") as f:
|
||||
json.dump(kp_dict, f, indent=1)
|
||||
|
||||
|
||||
def export_shot_changes(res_path, target_path):
|
||||
"""
|
||||
Exports the phalp output of shot changes into a JSON file
|
||||
dict with the shot index for each frame in the sequence
|
||||
"""
|
||||
os.makedirs(os.path.dirname(target_path), exist_ok=True)
|
||||
print(f"exporting shot changes to {target_path}")
|
||||
|
||||
tracklet_data = joblib.load(res_path)
|
||||
frames = sorted(tracklet_data.keys())
|
||||
# print(frames)
|
||||
# print(tracklet_data[frames[0]].keys())
|
||||
# for frame in frames:
|
||||
# print(frame, tracklet_data[frame]["shot"])
|
||||
shots = np.cumsum([tracklet_data[frame]["shot"] for frame in frames])
|
||||
shots = shots.astype(int).tolist()
|
||||
shot_dict = {frame.split("/")[-1]: shot for frame, shot in zip(frames, shots)}
|
||||
with open(target_path, "w") as f:
|
||||
json.dump(shot_dict, f, indent=1)
|
||||
|
||||
|
||||
def unpack_frame(track_data):
|
||||
"""
|
||||
Unpack one frame of track data from the phalp output into JSON files for each track.
|
||||
Track data contains body pose, global translation and rotation,
|
||||
and keypoints for all tracks in the frame.
|
||||
Each track is unpacked into folders with a JSON for each frame.
|
||||
"""
|
||||
out_dict = {}
|
||||
for i, tid in enumerate(track_data["tid"]):
|
||||
if tid not in track_data["tracked_ids"]:
|
||||
continue
|
||||
|
||||
cam_trans = track_data["cam_trans"][i].squeeze() # (3,)
|
||||
|
||||
mano = track_data["mano"][i]
|
||||
betas = mano["betas"].squeeze() # (10,)
|
||||
hand_pose = mano["hand_pose"].squeeze() # (23, 3, 3)
|
||||
is_right = mano["is_right"]
|
||||
print(is_right, tid, track_data["tid"], track_data["tracked_ids"])
|
||||
assert is_right == tid, f'{is_right}, {tid}'
|
||||
hand_pose_aa = np.stack(
|
||||
[cv2.Rodrigues(x)[0].squeeze() for x in hand_pose], axis=0
|
||||
) # (23, 3)
|
||||
global_orient = mano["global_orient"].squeeze() # (3, 3)
|
||||
global_orient_aa = cv2.Rodrigues(global_orient)[0].squeeze() # (3,)
|
||||
|
||||
out_dict[tid] = {
|
||||
"betas": betas.tolist(),
|
||||
"body_pose": hand_pose_aa.tolist(),
|
||||
"global_orient": global_orient_aa.tolist(),
|
||||
"cam_trans": cam_trans.tolist(),
|
||||
"is_right": is_right
|
||||
}
|
||||
|
||||
return out_dict
|
||||
|
||||
|
||||
def export_sequence_results(hamer_res_path, track_dir, shot_path):
|
||||
print('exporting sequence results in export_hamer.py...')
|
||||
export_hamer_predictions(hamer_res_path, track_dir)
|
||||
export_vitpose_keypoints(hamer_res_path, track_dir)
|
||||
export_shot_changes(hamer_res_path, shot_path)
|
||||
@@ -0,0 +1,206 @@
|
||||
import argparse
|
||||
import imageio
|
||||
import os
|
||||
import subprocess
|
||||
|
||||
|
||||
def split_frame(videopath,
|
||||
out_dir,
|
||||
fps=30,
|
||||
ext="jpg",
|
||||
down_scale=1,
|
||||
start_sec=0,
|
||||
end_sec=-1,
|
||||
overwrite=False,
|
||||
**kwargs,):
|
||||
import cv2
|
||||
base_name = videopath.replace('.mp4', '').replace('.MP4', '').replace('.avi', '')
|
||||
if os.path.isfile(videopath):
|
||||
print(f'processing {videopath} .....')
|
||||
# filepath = os.path.dirname(videopath)
|
||||
# videoname = os.path.basename(videopath)[:-4]
|
||||
if not os.path.exists(videopath):
|
||||
raise('{} not exists!'.format(videopath))
|
||||
|
||||
vc = cv2.VideoCapture(videopath)
|
||||
|
||||
path = out_dir
|
||||
if not os.path.exists(path):
|
||||
os.makedirs(path)
|
||||
|
||||
if vc.isOpened():
|
||||
rval , frame = vc.read()
|
||||
if '.MOV' not in videopath:
|
||||
cv2.imencode('.jpg',frame)[1].tofile(path +'/' + str(1).zfill(6) + '.jpg')
|
||||
else:
|
||||
cv2.imencode('.jpg',frame[::-1, ::-1])[1].tofile(path +'/' + str(0).zfill(6) + '.jpg')
|
||||
else:
|
||||
rval = False
|
||||
|
||||
Frame = 2
|
||||
timeF = 1
|
||||
while rval:
|
||||
rval,frame = vc.read()
|
||||
if rval==False:
|
||||
break
|
||||
|
||||
if '.MOV' not in videopath:
|
||||
cv2.imencode('.jpg',frame)[1].tofile(path +'/' + str(Frame).zfill(6) + '.jpg')
|
||||
else:
|
||||
cv2.imencode('.jpg',frame[::-1, ::-1])[1].tofile(path +'/' + str(Frame).zfill(6) + '.jpg')
|
||||
|
||||
Frame += 1
|
||||
elif os.path.isdir(videopath):
|
||||
print()
|
||||
print('moving ', base_name , 'to ', out_dir)
|
||||
os.system(f"cp -r {videopath} {out_dir}")
|
||||
raise ValueError
|
||||
return 0
|
||||
|
||||
else:
|
||||
print(os.path.abspath(videopath), os.path.exists(videopath))
|
||||
raise ValueError(f"{videopath} does not exist or is not a valid file or directory.")
|
||||
|
||||
return 0
|
||||
|
||||
def video_to_frames(
|
||||
path,
|
||||
out_dir,
|
||||
fps=30,
|
||||
ext="jpg",
|
||||
down_scale=1,
|
||||
start_sec=0,
|
||||
end_sec=-1,
|
||||
overwrite=False,
|
||||
**kwargs,
|
||||
):
|
||||
"""
|
||||
:param path
|
||||
:param out_dir
|
||||
:param fps
|
||||
:param down_scale (optional int)
|
||||
"""
|
||||
path = str(path)
|
||||
base_name = path.replace('.mp4', '').replace('.MP4', '').replace('.avi', '')
|
||||
print(path, base_name)
|
||||
print(os.path.isfile(path), os.path.isdir(base_name))
|
||||
|
||||
if os.path.isfile(path):
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
|
||||
arg_str = f"-copyts -qscale:v 2 -vf fps={fps}"
|
||||
if down_scale != 1:
|
||||
arg_str = f"{arg_str},scale='iw/{down_scale}:ih/{down_scale}'"
|
||||
if start_sec > 0:
|
||||
arg_str = f"{arg_str} -ss {start_sec}"
|
||||
if end_sec > start_sec:
|
||||
arg_str = f"{arg_str} -to {end_sec}"
|
||||
|
||||
yn = "-y" if overwrite else "-n"
|
||||
cmd = f"ffmpeg -i {path} {arg_str} {out_dir}/%06d.{ext} {yn}"
|
||||
print(cmd)
|
||||
elif os.path.isdir(base_name):
|
||||
os.system(f"cp -r {base_name} {out_dir}")
|
||||
return 0
|
||||
else:
|
||||
print(path, os.path.exists(path))
|
||||
raise ValueError
|
||||
|
||||
return subprocess.call(cmd, shell=True, stdin=subprocess.PIPE)
|
||||
|
||||
# def video_to_frames(
|
||||
# path,
|
||||
# out_dir,
|
||||
# fps=30,
|
||||
# ext="jpg",
|
||||
# down_scale=1,
|
||||
# start_sec=0,
|
||||
# end_sec=-1,
|
||||
# overwrite=False,
|
||||
# **kwargs,
|
||||
# ):
|
||||
# """
|
||||
# :param path
|
||||
# :param out_dir
|
||||
# :param fps
|
||||
# :param down_scale (optional int)
|
||||
# """
|
||||
# path = str(path)
|
||||
# base_name = path.replace('.mp4', '').replace('.MP4', '').replace('.avi', '')
|
||||
# print(path, base_name)
|
||||
# print(os.path.isfile(path), os.path.isdir(base_name))
|
||||
|
||||
# if os.path.isfile(path):
|
||||
# os.makedirs(out_dir, exist_ok=True)
|
||||
|
||||
# arg_str = f"-copyts -qscale:v 2 -vf fps={fps}"
|
||||
# if down_scale != 1:
|
||||
# arg_str = f"{arg_str},scale='iw/{down_scale}:ih/{down_scale}'"
|
||||
# if start_sec > 0:
|
||||
# arg_str = f"{arg_str} -ss {start_sec}"
|
||||
# if end_sec > start_sec:
|
||||
# arg_str = f"{arg_str} -to {end_sec}"
|
||||
|
||||
# yn = "-y" if overwrite else "-n"
|
||||
# # Changed %06d to start from 0 instead of 1
|
||||
# cmd = f"ffmpeg -i {path} {arg_str} -start_number 0 {out_dir}/%06d.{ext} {yn}"
|
||||
# print(cmd)
|
||||
# elif os.path.isdir(base_name):
|
||||
# os.system(f"cp -r {base_name} {out_dir}")
|
||||
# return 0
|
||||
# else:
|
||||
# print(path, os.path.exists(path))
|
||||
# raise ValueError
|
||||
|
||||
# return subprocess.call(cmd, shell=True, stdin=subprocess.PIPE)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--data_root", type=str, required=True, help="path to video")
|
||||
parser.add_argument(
|
||||
"--out_root", type=str, required=True, help="output dir for frames"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--seqs",
|
||||
nargs="*",
|
||||
default=None,
|
||||
help="[optional] sequences to run, default runs all available",
|
||||
)
|
||||
parser.add_argument("--fps", type=int, default=30, help="fps to extract frames")
|
||||
parser.add_argument(
|
||||
"--ext", type=str, default="jpg", help="output filetype for frames"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--down_scale", type=int, default=1, help="scale to extract frames"
|
||||
)
|
||||
parser.add_argument(
|
||||
"-ss", "--start_sec", type=float, default=0, help="seconds to start_sec"
|
||||
)
|
||||
parser.add_argument(
|
||||
"-es", "--end_sec", type=float, default=-1, help="seconds to end_sec"
|
||||
)
|
||||
parser.add_argument(
|
||||
"-y", "--overwrite", action="store_true", help="overwrite if already exist"
|
||||
)
|
||||
args = parser.parse_args()
|
||||
seqs_all = os.listdir(args.data_root)
|
||||
if args.seqs is None:
|
||||
args.seqs = seqs_all
|
||||
|
||||
for seq in args.seqs:
|
||||
path = os.path.join(args.data_root, seq)
|
||||
print(f"EXTRACTING FRAMES FROM {path}")
|
||||
assert os.path.isfile(path)
|
||||
seq_name = os.path.splitext(os.path.basename(path.rstrip("/")))[0]
|
||||
out_dir = os.path.join(args.out_root, seq_name)
|
||||
video_to_frames(
|
||||
path,
|
||||
out_dir,
|
||||
args.fps,
|
||||
args.ext,
|
||||
args.down_scale,
|
||||
args.start_sec,
|
||||
args.end_sec,
|
||||
args.overwrite,
|
||||
)
|
||||
@@ -0,0 +1,122 @@
|
||||
import os
|
||||
import subprocess
|
||||
import multiprocessing as mp
|
||||
from concurrent import futures
|
||||
|
||||
from preproc.datasets import update_args
|
||||
from preproc.export_hamer import export_sequence_results
|
||||
|
||||
ROOT_DIR = os.path.abspath(f"{__file__}/../../../")
|
||||
SRC_DIR = os.path.join(ROOT_DIR, "third-party/hamer")
|
||||
|
||||
def launch_hamer(gpus, seq, img_dir, res_dir, name, datatype, overwrite=False):
|
||||
"""
|
||||
run hamer using GPU pool
|
||||
"""
|
||||
cur_proc = mp.current_process()
|
||||
print("PROCESS", cur_proc.name, cur_proc._identity)
|
||||
# 1-indexed processes
|
||||
# worker_id = cur_proc._identity[0] - 1 if len(cur_proc._identity) > 0 else 0
|
||||
# gpu = gpus[worker_id % len(gpus)]
|
||||
gpu = gpus[0]
|
||||
|
||||
HAMER_DIR = SRC_DIR
|
||||
print("HAMER DIR", HAMER_DIR)
|
||||
|
||||
cmd_args = [
|
||||
f"cd {HAMER_DIR};",
|
||||
f"CUDA_VISIBLE_DEVICES={gpu}",
|
||||
"python -u run.py",
|
||||
f"--img_folder {img_dir} ",
|
||||
f"--res_folder {res_dir}/demo_{name}.pkl ",
|
||||
f"--batch_size=48 --side_view --save_mesh --full_frame",
|
||||
f"--type {datatype}",
|
||||
f"--checkpoint {ROOT_DIR}",
|
||||
"--render"
|
||||
]
|
||||
|
||||
cmd = " ".join(cmd_args)
|
||||
print(cmd)
|
||||
return subprocess.call(cmd, shell=True)
|
||||
|
||||
|
||||
def process_seq(
|
||||
gpus,
|
||||
out_root,
|
||||
seq,
|
||||
img_dir,
|
||||
out_name="hamer_out",
|
||||
datatype=None,
|
||||
track_name="track_preds",
|
||||
shot_name="shot_idcs",
|
||||
overwrite=False,
|
||||
):
|
||||
"""
|
||||
Run and export HAMER results
|
||||
"""
|
||||
name = os.path.basename(seq)
|
||||
res_root = f"{out_root}/{out_name}/{seq}"
|
||||
os.makedirs(res_root, exist_ok=True)
|
||||
res_dir = os.path.join(res_root, "results")
|
||||
res_path = f"{res_root}/{name}.pkl"
|
||||
|
||||
if overwrite or not os.path.isfile(res_path):
|
||||
res = launch_hamer(gpus, seq, img_dir, res_dir, name, datatype, overwrite)
|
||||
print(f'rename {res_dir}/demo_{name}.pkl into ', res_path)
|
||||
os.rename(f"{res_dir}/demo_{name}.pkl", res_path)
|
||||
assert res == 0, "HAMER FAILED"
|
||||
|
||||
# export the HAMER predictions
|
||||
track_dir = f"{out_root}/{track_name}/{seq}"
|
||||
shot_path = f"{out_root}/{shot_name}/{seq}.json"
|
||||
|
||||
export_sequence_results(res_path, track_dir, shot_path)
|
||||
return 0
|
||||
|
||||
|
||||
def get_out_dir(src_root, src_dir, src_token, out_token):
|
||||
"""
|
||||
:param src_root (str) root of all data
|
||||
:param src_dir (str) img input dir
|
||||
:param src_token (str) parent name of image input dir
|
||||
:param out_token (str) name of output dir
|
||||
"""
|
||||
src_suffix = src_dir.removeprefix(src_root)
|
||||
out_dir = f"{out_root}/{src_suffix}"
|
||||
return out_dir.replace(src_token, out_token)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--type", default="posetrack", help="dataset to process")
|
||||
parser.add_argument("--root", default=None, help="root dir of data, default None")
|
||||
parser.add_argument("--split", default="val", help="split of dataset, default val")
|
||||
parser.add_argument(
|
||||
"--img_name", default=None, help="input image directory name, default None"
|
||||
)
|
||||
parser.add_argument("--seqs", nargs="*", default=None)
|
||||
parser.add_argument("--gpus", nargs="*", default=[0])
|
||||
parser.add_argument("-y", "--overwrite", action="store_true")
|
||||
|
||||
args = parser.parse_args()
|
||||
args = update_args(args)
|
||||
|
||||
out_root = f"{args.root}/slahmr/{args.split}"
|
||||
|
||||
print(f"running phalp on {len(args.img_dirs)} image directories")
|
||||
if len(args.gpus) > 1:
|
||||
with futures.ProcessPoolExecutor(max_workers=len(args.gpus)) as exe:
|
||||
for img_dir, seq in zip(args.img_dirs, args.seqs):
|
||||
exe.submit(
|
||||
process_seq,
|
||||
args.gpus,
|
||||
out_root,
|
||||
seq,
|
||||
img_dir,
|
||||
overwrite=args.overwrite,
|
||||
)
|
||||
else:
|
||||
for img_dir, seq in zip(args.img_dirs, args.seqs):
|
||||
process_seq(args.gpus, out_root, seq, img_dir, overwrite=args.overwrite)
|
||||
@@ -0,0 +1,186 @@
|
||||
import os
|
||||
import glob
|
||||
import json
|
||||
import numpy as np
|
||||
|
||||
import subprocess
|
||||
from concurrent import futures
|
||||
import multiprocessing as mp
|
||||
|
||||
from preproc.datasets import update_args
|
||||
|
||||
|
||||
SRC_DIR = os.path.abspath(f"{__file__}/../")
|
||||
|
||||
|
||||
def isimage(path):
|
||||
ext = os.path.splitext(path)[-1].lower()
|
||||
return ext == ".png" or ext == ".jpg" or ext == ".jpeg"
|
||||
|
||||
|
||||
def split_frames_equal(img_dir, seq_len=-1):
|
||||
"""
|
||||
split the sequence into subsequences of seq_len
|
||||
returns start and end indices for each subsequence
|
||||
"""
|
||||
if seq_len == -1: # don't split
|
||||
return [(0, -1)]
|
||||
|
||||
image_list = sorted(list(filter(isimage, os.listdir(img_dir))))
|
||||
num_imgs = len(image_list)
|
||||
splits = list(range(0, num_imgs, seq_len))
|
||||
splits[-1] = num_imgs # add remainder to last subseq
|
||||
return list(zip(splits[:-1], splits[1:]))
|
||||
|
||||
|
||||
def split_frames_shots(img_dir, shot_path, pad_shot=0, min_len=0):
|
||||
"""
|
||||
split the sequence into subsequences according by shots in the video
|
||||
returns start and end indices for each subsequence
|
||||
"""
|
||||
if shot_path is None or not os.path.isfile(shot_path):
|
||||
print(f"{shot_path} DOES NOT EXIST, USING WHOLE SEQUENCE")
|
||||
return [(0, -1)], [0]
|
||||
|
||||
image_list = sorted(list(filter(isimage, os.listdir(img_dir))))#[:500]
|
||||
num_frames = len(image_list)
|
||||
|
||||
with open(shot_path, "r") as f:
|
||||
shot_dict = json.load(f)
|
||||
|
||||
# print(shot_dict.keys())
|
||||
shot_idcs = np.array([shot_dict[name] for name in image_list])
|
||||
subseqs = []
|
||||
for i in np.unique(shot_idcs):
|
||||
idcs = np.where(shot_idcs == i)[0]
|
||||
start, end = idcs.min(), idcs.max() + 1
|
||||
if pad_shot > 0:
|
||||
if start > 0:
|
||||
start = start + pad_shot
|
||||
if end < num_frames:
|
||||
end = end - pad_shot
|
||||
if (end - start) < min_len: # this shot is too short
|
||||
continue
|
||||
subseqs.append((start, end))
|
||||
return subseqs, np.unique(shot_idcs)
|
||||
|
||||
|
||||
def split_sequence(args, img_dir, seq):
|
||||
subseqs = split_frames_equal(img_dir, args.seq_len)
|
||||
return subseqs, [None for _ in subseqs]
|
||||
|
||||
|
||||
def get_out_dir(args, seq, shot_idx=None, start=0, end=-1):
|
||||
if "egobody" in args.type or "3dpw" in args.type:
|
||||
if args.use_intrins:
|
||||
out_name = "cameras_intrins"
|
||||
else:
|
||||
out_name = "cameras_default"
|
||||
else:
|
||||
out_name = "cameras"
|
||||
|
||||
if args.seq_len > 0:
|
||||
out_name = f"{out_name}_split"
|
||||
|
||||
out_dir = f"{args.root}/slahmr/{args.split}/{out_name}/{seq}"
|
||||
|
||||
if args.seq_len > 0:
|
||||
return f"{out_dir}/{start}-{end}"
|
||||
|
||||
if args.type == "posetrack" and shot_idx is not None:
|
||||
return f"{out_dir}/shot-{shot_idx}"
|
||||
|
||||
return out_dir
|
||||
|
||||
|
||||
def get_intrins_path(data_type, data_root, seq):
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
def get_command(img_dir, out_dir, start=0, end=-1, intrins_path=None, overwrite=False):
|
||||
cmd_args = [
|
||||
f"python {SRC_DIR}/run_slam.py",
|
||||
"-i",
|
||||
img_dir,
|
||||
"--map_dir",
|
||||
out_dir,
|
||||
"--disable_vis",
|
||||
]
|
||||
if start != 0:
|
||||
cmd_args += ["--start", str(start)]
|
||||
if end != -1:
|
||||
cmd_args += ["--end", str(end)]
|
||||
|
||||
if intrins_path is not None:
|
||||
cmd_args += ["--intrins_path", intrins_path]
|
||||
|
||||
if overwrite:
|
||||
cmd_args.append("--overwrite")
|
||||
|
||||
cmd = " ".join(cmd_args)
|
||||
return cmd
|
||||
|
||||
|
||||
def check_intrins(data_type, data_root, intrins_path, img_dir):
|
||||
assert intrins_path is not None
|
||||
out_name = intrins_path.split(data_root)[1].strip("/").split("/")[0]
|
||||
out_root = f"{data_root}/{out_name}"
|
||||
return intrins_path
|
||||
|
||||
|
||||
def get_slam_command(args, img_dir, seq, shot_idx=None, start=0, end=-1):
|
||||
out_dir = get_out_dir(args, seq, shot_idx, start, end)
|
||||
intrins_path = None
|
||||
if args.use_intrins:
|
||||
intrins_path = get_intrins_path(args.type, args.root, seq)
|
||||
if intrins_path is not None:
|
||||
intrins_path = check_intrins(args.type, args.root, intrins_path, img_dir)
|
||||
return get_command(
|
||||
img_dir,
|
||||
out_dir,
|
||||
start=start,
|
||||
end=end,
|
||||
intrins_path=intrins_path,
|
||||
overwrite=args.overwrite,
|
||||
)
|
||||
|
||||
|
||||
def launch_job(gpus, cmd):
|
||||
cur_proc = mp.current_process()
|
||||
print(cur_proc.name, cur_proc._identity)
|
||||
worker_id = cur_proc._identity[0] - 1
|
||||
gpu = gpus[worker_id % len(gpus)]
|
||||
cmd = f"CUDA_VISIBLE_DEVICES={gpu} {cmd}"
|
||||
print(cmd)
|
||||
subprocess.call(cmd, shell=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--type", default="posetrack", help="dataset to process")
|
||||
parser.add_argument(
|
||||
"--root", default=None, help="root of data to process, default None"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--split", default="val", help="split of dataset to process, default val"
|
||||
)
|
||||
parser.add_argument("--gpus", nargs="*", default=[0])
|
||||
parser.add_argument("--seqs", nargs="*", default=None)
|
||||
parser.add_argument("--seq_len", type=int, default=-1)
|
||||
parser.add_argument(
|
||||
"--use_intrins", action="store_true", help="use GT intrinsics if available"
|
||||
)
|
||||
parser.add_argument("-y", "--overwrite", action="store_true")
|
||||
args = parser.parse_args()
|
||||
args = update_args(args)
|
||||
|
||||
print(f"Running SLAM on {len(args.seqs)} sequences")
|
||||
with futures.ProcessPoolExecutor(max_workers=len(args.gpus)) as ex:
|
||||
for img_dir, seq in zip(args.img_dirs, args.seqs):
|
||||
subseqs, shot_idcs = split_sequence(args, img_dir, seq)
|
||||
for shot, (start, end) in zip(shot_idcs, subseqs):
|
||||
cmd = get_slam_command(args, img_dir, seq, shot, start, end)
|
||||
print(cmd)
|
||||
ex.submit(launch_job, args.gpus, cmd)
|
||||
+414
@@ -0,0 +1,414 @@
|
||||
import os
|
||||
|
||||
ROOT_DIR = os.path.abspath(f"{__file__}/../../../")
|
||||
SRC_DIR = os.path.join(ROOT_DIR, "third-party/DROID-SLAM")
|
||||
print("PROJ SRC", ROOT_DIR)
|
||||
print("DROID SRC", SRC_DIR)
|
||||
|
||||
import sys
|
||||
|
||||
sys.path.append(SRC_DIR)
|
||||
sys.path.append(f"{SRC_DIR}/droid_slam")
|
||||
|
||||
import glob
|
||||
import shutil
|
||||
import argparse
|
||||
from tqdm import tqdm
|
||||
|
||||
import json
|
||||
|
||||
import cv2
|
||||
import trimesh
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from lietorch import SE3
|
||||
from droid import Droid
|
||||
import droid_backends
|
||||
|
||||
|
||||
def get_image(image):
|
||||
return image.permute(1, 2, 0).cpu().numpy()[..., ::-1]
|
||||
|
||||
|
||||
def show_image(image):
|
||||
image = image.permute(1, 2, 0).cpu().numpy()
|
||||
cv2.imshow("image", image / 255.0)
|
||||
cv2.waitKey(1)
|
||||
|
||||
|
||||
def get_hwf(test_path):
|
||||
test_img = cv2.imread(test_path)
|
||||
H, W, _ = test_img.shape
|
||||
F = 0.5 * (H + W)
|
||||
return H, W, F
|
||||
|
||||
|
||||
def isimage(path):
|
||||
ext = os.path.splitext(path)[-1].lower()
|
||||
return ext == ".png" or ext == ".jpg" or ext == ".jpeg"
|
||||
|
||||
|
||||
def get_image_files(img_dir, stride, start=0, end=-1):
|
||||
image_list = sorted(list(filter(isimage, os.listdir(img_dir))))
|
||||
end = len(image_list) + 1 + end if end < 0 else end
|
||||
return [os.path.join(img_dir, name) for name in image_list[start:end:stride]]
|
||||
|
||||
|
||||
def load_intrins(intrins_path, image_files, start=0, end=-1):
|
||||
N = len(image_files)
|
||||
H, W, F = get_hwf(image_files[0])
|
||||
# intrinsics for default FOV
|
||||
default = torch.tensor([F, F, W / 2, H / 2, W, H])[None].repeat(N, 1)
|
||||
if intrins_path is None:
|
||||
return default
|
||||
assert os.path.isfile(intrins_path), f"{intrins_path} DOES NOT EXIST"
|
||||
assert intrins_path.endswith(".txt"), f"{intrins_path} INCORRECT EXT"
|
||||
|
||||
print(f"LOADING INTRINSICS FROM {intrins_path}")
|
||||
intrins = torch.from_numpy(np.loadtxt(intrins_path).astype(np.float32))
|
||||
print(intrins.shape, N)
|
||||
if intrins.ndim < 2: # assume intrins is same for all frames
|
||||
intrins = intrins[None].repeat(N, 1)
|
||||
else:
|
||||
end = N + 1 + end if end < 0 else end
|
||||
intrins = intrins[start:end]
|
||||
print(intrins.shape)
|
||||
if intrins.shape[1] == 4: # fx, fy, cx, cy, adding on W and H
|
||||
ones = torch.ones_like(intrins[:, :1])
|
||||
intrins = torch.cat([intrins, ones * W, ones * H], dim=-1)
|
||||
return intrins
|
||||
|
||||
|
||||
def image_stream(image_files, intrins_all):
|
||||
"""image generator"""
|
||||
|
||||
N = len(image_files)
|
||||
|
||||
assert intrins_all.shape[0] == N
|
||||
assert intrins_all.shape[1] >= 4
|
||||
|
||||
for t, imfile in enumerate(image_files):
|
||||
image = cv2.imread(imfile)
|
||||
if image is None:
|
||||
print(imfile, "is none, exiting")
|
||||
sys.exit(1)
|
||||
|
||||
h0, w0, _ = image.shape
|
||||
h1 = int(h0 * np.sqrt((384 * 512) / (h0 * w0)))
|
||||
w1 = int(w0 * np.sqrt((384 * 512) / (h0 * w0)))
|
||||
|
||||
image = cv2.resize(image, (w1, h1))
|
||||
image = image[: h1 - h1 % 8, : w1 - w1 % 8]
|
||||
image = torch.as_tensor(image).permute(2, 0, 1)
|
||||
|
||||
fx, fy, cx, cy = intrins_all[t, :4]
|
||||
sx, sy = w1 / w0, h1 / h0
|
||||
frame_intrins = torch.as_tensor([fx * sx, fy * sy, cx * sx, cy * sy])
|
||||
yield t, image[None], frame_intrins
|
||||
|
||||
|
||||
def unpack_video(video, depth_filter_thresh=0.005):
|
||||
"""
|
||||
extract the video of keyframe disps and poses
|
||||
returns:
|
||||
(T) timestamps
|
||||
(T, 7) world2cam pose for each keyframe
|
||||
(T, H, W) disps
|
||||
(T, H, W, 3) images
|
||||
(T, H, W) mask of valid pixels
|
||||
(4) intrinsics (fx, fy, cx, cy)
|
||||
"""
|
||||
with torch.no_grad():
|
||||
with video.get_lock():
|
||||
# dirty frames have disparity normalized over all previous frames
|
||||
# should probably change this in visualizer
|
||||
t = video.counter.value
|
||||
|
||||
print(f"{t} keyframes in map")
|
||||
|
||||
# normalization happens after all frames have been processed
|
||||
if t < 1:
|
||||
print("no keyframes")
|
||||
return
|
||||
|
||||
# taken from visualization.py#80
|
||||
tstamps = video.tstamp[:t]
|
||||
dirty_index = torch.arange(t, device=video.poses.device)
|
||||
|
||||
poses = torch.index_select(video.poses, 0, dirty_index) # (T, 7)
|
||||
disps = torch.index_select(video.disps, 0, dirty_index) # (T, H, W)
|
||||
images = torch.index_select(video.images, 0, dirty_index) # (T, 3, H, W)
|
||||
|
||||
images = images[:, [2, 1, 0], 3::8, 3::8].permute(0, 2, 3, 1) / 255.0
|
||||
intrins = video.intrinsics[0] # (4,)
|
||||
|
||||
# get mask of static points
|
||||
thresh = depth_filter_thresh * torch.ones_like(disps.mean(dim=[1, 2]))
|
||||
count = droid_backends.depth_filter(
|
||||
video.poses, video.disps, intrins, dirty_index, thresh
|
||||
)
|
||||
disp_mean = disps.mean(dim=(1, 2), keepdim=True).clamp_min(0.2)
|
||||
valid = (count >= 2) & (disps > 0.5 * disp_mean) # (T, H, W)
|
||||
return {
|
||||
"tstamps": tstamps.cpu(),
|
||||
"w2c": poses.cpu(),
|
||||
"disps": disps.cpu(),
|
||||
"images": images.cpu(),
|
||||
"valid": valid.cpu(),
|
||||
"intrins": intrins.cpu(),
|
||||
}
|
||||
|
||||
|
||||
def get_keyframe_map(video_dict):
|
||||
"""
|
||||
:param video_dict (dict) unpacked droid slam video
|
||||
:returns
|
||||
(T, 4, 4) camera to world matrix
|
||||
(N, 3) valid points in world
|
||||
(N, 3) colors of valid points
|
||||
"""
|
||||
poses = video_dict["w2c"]
|
||||
images = video_dict["images"]
|
||||
valid = video_dict["valid"]
|
||||
t = len(poses)
|
||||
c2w = SE3(poses).inv()
|
||||
|
||||
# `iproj` only has a cuda backend; if we pass in CPU tensors, it will silently
|
||||
# return zeros
|
||||
points = droid_backends.iproj(
|
||||
c2w.data.cuda(),
|
||||
video_dict["disps"].cuda(),
|
||||
video_dict["intrins"].cuda(),
|
||||
)
|
||||
valid_pts = torch.cat([points[i, valid[i]] for i in range(t)], dim=0)
|
||||
valid_rgb = torch.cat([images[i, valid[i]] for i in range(t)], dim=0)
|
||||
return (
|
||||
c2w.matrix().detach().cpu(),
|
||||
valid_pts.detach().cpu(),
|
||||
valid_rgb.detach().cpu(),
|
||||
)
|
||||
|
||||
|
||||
def get_frame_cameras(droid, img_paths, intrins_all):
|
||||
N = len(img_paths)
|
||||
|
||||
t = droid.video.counter.value
|
||||
print(f"{t} keyframes in map")
|
||||
|
||||
if t > 1:
|
||||
with torch.no_grad():
|
||||
# localize all frames and get edges into keyframe graph
|
||||
# returns 7D tensor (3D trans, 4D quat)
|
||||
c2w = droid.terminate(image_stream(img_paths, intrins_all))
|
||||
c2w = torch.from_numpy(c2w.astype(np.float32))
|
||||
return SE3(c2w).inv().matrix()
|
||||
|
||||
return torch.eye(4)[None].repeat(N, 1, 1)
|
||||
|
||||
|
||||
def save_camera_json(path, extrins, intrins, yup=False):
|
||||
"""
|
||||
:param path
|
||||
:param extrins (N, 4, 4)
|
||||
:param intrins (N, 4)
|
||||
"""
|
||||
N = extrins.shape[0]
|
||||
if yup:
|
||||
T = torch.tensor(
|
||||
[[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]],
|
||||
dtype=torch.float32,
|
||||
)
|
||||
extrins = torch.matmul(T[None], extrins)
|
||||
with open(path, "w") as f:
|
||||
json.dump(
|
||||
{
|
||||
"rotation": extrins[:, :3, :3].reshape((N, 9)).tolist(),
|
||||
"translation": extrins[:, :3, 3].tolist(),
|
||||
"intrinsics": intrins.tolist(),
|
||||
},
|
||||
f,
|
||||
indent=1,
|
||||
)
|
||||
|
||||
|
||||
def save_pcl(path, points, colors, yup=False):
|
||||
assert len(points) == len(colors)
|
||||
if yup:
|
||||
T = torch.tensor(
|
||||
[[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]],
|
||||
dtype=torch.float32,
|
||||
)
|
||||
points = torch.einsum("ij,...j->...i", T, points) # (*, 3)
|
||||
points = points.reshape(-1, 3).numpy()
|
||||
colors = colors.reshape(-1, 3).numpy()
|
||||
obj = trimesh.Trimesh(vertices=points, vertex_colors=colors)
|
||||
with open(path, "wb") as f:
|
||||
f.write(trimesh.exchange.ply.export_ply(obj))
|
||||
|
||||
|
||||
def save_keyframe_map(out_dir, kf_nodes):
|
||||
"""
|
||||
Saves everything with x right y up z back
|
||||
(converts from x right y down z forward)
|
||||
:param out_dir
|
||||
:param kf_nodes (dict)
|
||||
"""
|
||||
c2w, points, colors = get_keyframe_map(kf_nodes)
|
||||
intrins = kf_nodes["intrins"].detach().cpu()
|
||||
|
||||
print("c2w", c2w.shape)
|
||||
print("points", points.shape)
|
||||
print("colors", colors.shape)
|
||||
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
save_pcl(f"{out_dir}/map_points.ply", points, colors)
|
||||
save_camera_json(f"{out_dir}/map_cameras.json", c2w, intrins)
|
||||
print(f"saved {len(c2w)} keyframes to {out_dir}")
|
||||
|
||||
|
||||
def save_cameras(map_dir, frame_w2c, intrins):
|
||||
os.makedirs(map_dir, exist_ok=True)
|
||||
print("SAVING CAMERA INTRINS AND EXTRINS IN", map_dir)
|
||||
|
||||
# save all camera extrinsics and intrinsic
|
||||
W, H = intrins[0, 4], intrins[0, 5]
|
||||
focal = intrins[:, :2].mean()
|
||||
print("INTRINS", intrins[0])
|
||||
np.savez(
|
||||
f"{map_dir}/cameras.npz",
|
||||
height=H,
|
||||
width=W,
|
||||
focal=focal,
|
||||
intrins=intrins[:, :4],
|
||||
w2c=frame_w2c,
|
||||
)
|
||||
|
||||
# save visualization
|
||||
frame_intrins = torch.tensor([focal, focal, W / 2, H / 2, W, H])
|
||||
frame_c2w = torch.linalg.inv(frame_w2c)
|
||||
save_camera_json(f"{map_dir}/frame_cameras.json", frame_c2w, intrins)
|
||||
|
||||
|
||||
def main(args):
|
||||
args.stereo = False
|
||||
torch.multiprocessing.set_start_method("spawn")
|
||||
|
||||
droid = None
|
||||
|
||||
img_paths = get_image_files(
|
||||
args.img_dir, args.stride, start=args.start, end=args.end
|
||||
)
|
||||
if args.map_dir is not None:
|
||||
os.makedirs(args.map_dir, exist_ok=True)
|
||||
camera_file = os.path.join(args.map_dir, "cameras.npz")
|
||||
if os.path.isfile(camera_file) and not args.overwrite:
|
||||
print(glob.glob(f"{args.map_dir}/*.npz"), "already exist, skipping")
|
||||
return
|
||||
|
||||
intrins_all = load_intrins(
|
||||
args.intrins_path, img_paths, start=args.start, end=args.end
|
||||
)
|
||||
print("intrins shape", intrins_all.shape, "num images", len(img_paths))
|
||||
|
||||
for t, image, intrinsics in tqdm(image_stream(img_paths, intrins_all)):
|
||||
if t < args.t0:
|
||||
continue
|
||||
|
||||
if not args.disable_vis:
|
||||
show_image(image[0])
|
||||
|
||||
if droid is None:
|
||||
args.image_size = [image.shape[2], image.shape[3]]
|
||||
droid = Droid(args)
|
||||
|
||||
droid.track(t, image, intrinsics=intrinsics)
|
||||
|
||||
if args.map_dir is None:
|
||||
return
|
||||
|
||||
# save cameras
|
||||
frame_w2c = get_frame_cameras(droid, img_paths, intrins_all)
|
||||
save_cameras(args.map_dir, frame_w2c, intrins_all)
|
||||
|
||||
# save keyframe cameras and points
|
||||
kf_nodes = unpack_video(droid.video)
|
||||
save_keyframe_map(args.map_dir, kf_nodes)
|
||||
|
||||
|
||||
def get_slam_parser():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--stride", default=1, type=int, help="frame stride")
|
||||
parser.add_argument("--t0", default=0, type=int, help="starting frame")
|
||||
parser.add_argument("--weights", default=f"{ROOT_DIR}/_DATA/droid.pth")
|
||||
parser.add_argument("--buffer", type=int, default=512)
|
||||
parser.add_argument("--image_size", default=[240, 320])
|
||||
parser.add_argument("--disable_vis", action="store_true")
|
||||
|
||||
parser.add_argument(
|
||||
"--beta",
|
||||
type=float,
|
||||
default=0.3,
|
||||
help="weight for translation / rotation components of flow",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--filter_thresh",
|
||||
type=float,
|
||||
default=1.0,
|
||||
help="how much motion before considering new keyframe",
|
||||
)
|
||||
parser.add_argument("--warmup", type=int, default=8, help="number of warmup frames")
|
||||
parser.add_argument(
|
||||
"--keyframe_thresh",
|
||||
type=float,
|
||||
default=4.0,
|
||||
help="threshold to create a new keyframe",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--frontend_thresh",
|
||||
type=float,
|
||||
default=16.0,
|
||||
help="add edges between frames whithin this distance",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--frontend_window", type=int, default=25, help="frontend optimization window"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--frontend_radius",
|
||||
type=int,
|
||||
default=2,
|
||||
help="force edges between frames within radius",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--frontend_nms", type=int, default=1, help="non-maximal supression of edges"
|
||||
)
|
||||
|
||||
parser.add_argument("--backend_thresh", type=float, default=22.0)
|
||||
parser.add_argument("--backend_radius", type=int, default=2)
|
||||
parser.add_argument("--backend_nms", type=int, default=3)
|
||||
parser.add_argument("--upsample", action="store_true")
|
||||
return parser
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = get_slam_parser()
|
||||
parser.add_argument(
|
||||
"-i", "--img_dir", required=True, type=str, help="path to image directory"
|
||||
)
|
||||
parser.add_argument(
|
||||
"-o", "--map_dir", type=str, default=None, help="path to output directory"
|
||||
)
|
||||
parser.add_argument(
|
||||
"-y", "--overwrite", action="store_true", help="overwrite existing cameras"
|
||||
)
|
||||
parser.add_argument(
|
||||
"-g", "--save_graph", action="store_true", help="save the frame graph"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--intrins_path", default=None, help="path to camera intrinsics"
|
||||
)
|
||||
parser.add_argument("--start", default=0, type=int, help="start frame")
|
||||
parser.add_argument("--end", default=-1, type=int, help="end frame")
|
||||
|
||||
args = parser.parse_args()
|
||||
main(args)
|
||||
+204
@@ -0,0 +1,204 @@
|
||||
import os
|
||||
import glob
|
||||
import json
|
||||
import subprocess
|
||||
import random
|
||||
import numpy as np
|
||||
|
||||
import torch
|
||||
from torch.utils.data import DataLoader
|
||||
from torch.utils.tensorboard import SummaryWriter
|
||||
|
||||
from data import get_dataset_from_cfg, expand_source_paths
|
||||
|
||||
# from humor.humor_model import HumorModel
|
||||
from optim.base_scene import BaseSceneModel
|
||||
|
||||
from optim.optimizers import (
|
||||
RootOptimizer,
|
||||
SmoothOptimizer
|
||||
)
|
||||
from optim.output import (
|
||||
save_track_info,
|
||||
save_camera_json,
|
||||
save_input_poses,
|
||||
save_initial_predictions,
|
||||
)
|
||||
from vis.viewer import init_viewer
|
||||
from body_model import MANO
|
||||
from util.loaders import (
|
||||
# load_vposer,
|
||||
# load_state,
|
||||
# load_gmm,
|
||||
resolve_cfg_paths,
|
||||
)
|
||||
from util.logger import Logger
|
||||
from util.tensor import get_device, move_to, detach_all, to_torch
|
||||
|
||||
from run_vis import run_vis
|
||||
|
||||
import hydra
|
||||
from omegaconf import DictConfig, OmegaConf
|
||||
import time
|
||||
|
||||
N_STAGES = 3
|
||||
|
||||
import sys
|
||||
sys.path.append('src/human_body_prior')
|
||||
sys.path.append('HMP/')
|
||||
# print(sys.path)
|
||||
# print(torch.cuda.memory_summary())
|
||||
from HMP.fitting import run_prior
|
||||
|
||||
|
||||
def set_seed(seed=42):
|
||||
"""
|
||||
Set random seed for reproducibility
|
||||
"""
|
||||
random.seed(seed)
|
||||
np.random.seed(seed)
|
||||
torch.manual_seed(seed)
|
||||
torch.cuda.manual_seed(seed)
|
||||
torch.cuda.manual_seed_all(seed)
|
||||
torch.backends.cudnn.benchmark = False
|
||||
torch.backends.cudnn.deterministic = True
|
||||
os.environ['PYTHONHASHSEED'] = str(seed)
|
||||
|
||||
def run_opt(cfg, dataset, out_dir, device):
|
||||
a = time.time()
|
||||
args = cfg.data
|
||||
B = len(dataset)
|
||||
T = dataset.seq_len
|
||||
loader = DataLoader(dataset, batch_size=B, shuffle=False)
|
||||
|
||||
obs_data = move_to(next(iter(loader)), device)
|
||||
cam_data = move_to(dataset.get_camera_data(), device)
|
||||
print("Batch size (dataset_length), T (dataset.seq_len): ", B, len(dataset), T)
|
||||
print("OBS DATA", obs_data.keys())
|
||||
print("CAM DATA", cam_data.keys(), 'cam_R', cam_data['cam_R'])
|
||||
|
||||
# save cameras
|
||||
cam_R, cam_t = dataset.cam_data.cam2world()
|
||||
intrins = dataset.cam_data.intrins
|
||||
save_camera_json(f"cameras.json", cam_R, cam_t, intrins)
|
||||
|
||||
# check whether the cameras are static
|
||||
# if static, cannot optimize scale
|
||||
cfg.model.opt_scale &= not dataset.cam_data.is_static
|
||||
Logger.log(f"OPT SCALE {cfg.model.opt_scale}")
|
||||
|
||||
# loss weights for all stages
|
||||
all_loss_weights = cfg.optim.loss_weights
|
||||
assert all(len(wts) == N_STAGES for wts in all_loss_weights.values())
|
||||
stage_loss_weights = [
|
||||
{k: wts[i] for k, wts in all_loss_weights.items()} for i in range(N_STAGES)
|
||||
]
|
||||
max_loss_weights = {k: max(wts) for k, wts in all_loss_weights.items()}
|
||||
|
||||
# load models
|
||||
cfg = resolve_cfg_paths(cfg)
|
||||
cfg.paths.base_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "../"))
|
||||
paths = cfg.paths
|
||||
Logger.log(f"Loading pose prior from {paths.vposer}")
|
||||
# pose_prior, _ = load_model(paths.vposer, model_code=VPoser, remove_words_in_model_weights='vp_model.', disable_grad=True)
|
||||
# pose_prior = pose_prior.to(device)
|
||||
pose_prior = None
|
||||
|
||||
# load models
|
||||
cfg = resolve_cfg_paths(cfg)
|
||||
paths = cfg.paths
|
||||
Logger.log(f"Loading hand model")
|
||||
# Instantiate MANO model
|
||||
mano_cfg = {k.lower(): v for k,v in dict(cfg.MANO).items()}
|
||||
print('initializing MANO model with cfgs:', mano_cfg)
|
||||
hand_model = MANO(batch_size=B*T, pose2rot=True, **mano_cfg).to(device)
|
||||
|
||||
################################################################
|
||||
######################## optimization ##########################
|
||||
################################################################
|
||||
margs = cfg.model
|
||||
base_model = BaseSceneModel(
|
||||
B, T, hand_model, pose_prior, **margs
|
||||
)
|
||||
|
||||
base_model.initialize(obs_data, cam_data)
|
||||
base_model.to(device)
|
||||
|
||||
# save initial results for later visualization
|
||||
save_input_poses(dataset, os.path.join(out_dir, "hamer"), args.seq)
|
||||
save_initial_predictions(base_model, os.path.join(out_dir, "init"), args.seq)
|
||||
|
||||
opts = cfg.optim.options
|
||||
vis_scale = 0.25
|
||||
vis = None
|
||||
if opts.vis_every > 0:
|
||||
vis = init_viewer(
|
||||
dataset.img_size,
|
||||
cam_data["intrins"][0],
|
||||
vis_scale=vis_scale,
|
||||
bg_paths=dataset.sel_img_paths,
|
||||
fps=cfg.fps,
|
||||
)
|
||||
print("OPTIMIZER OPTIONS:", opts)
|
||||
|
||||
writer = SummaryWriter(out_dir)
|
||||
|
||||
print('start optimization')
|
||||
a = time.time()
|
||||
optim = RootOptimizer(base_model, stage_loss_weights, **opts)
|
||||
optim.run(obs_data, cfg.optim.root.num_iters, out_dir, vis, writer)
|
||||
|
||||
args = cfg.optim.smooth
|
||||
print(args)
|
||||
b = time.time()
|
||||
print('root optimization time: ', b - a)
|
||||
optim = SmoothOptimizer(
|
||||
base_model, stage_loss_weights, opt_scale=args.opt_scale, **opts
|
||||
)
|
||||
optim.run(obs_data, args.num_iters, out_dir, vis, writer)
|
||||
c = time.time()
|
||||
print('Smooth optimization time: ', c - b)
|
||||
|
||||
# HMP
|
||||
if cfg.run_prior and not glob.glob(os.path.join(out_dir, 'prior', '*_world_results.npz')):
|
||||
run_prior(cfg, dataset, out_dir, device, ['smooth_fit'], \
|
||||
obs_data, hand_model, cfg, cfg.data, os.path.join(out_dir, 'prior'))
|
||||
d = time.time()
|
||||
print('prior optimization time: ', d-c)
|
||||
|
||||
|
||||
@hydra.main(version_base=None, config_path="confs", config_name="config.yaml")
|
||||
def main(cfg: DictConfig):
|
||||
OmegaConf.register_new_resolver("eval", eval)
|
||||
print('run_opt.py: ', cfg)
|
||||
|
||||
# Set random seed
|
||||
set_seed(cfg.get('seed', 42))
|
||||
|
||||
out_dir = os.getcwd()
|
||||
print("out_dir", out_dir)
|
||||
Logger.init(f"{out_dir}/opt_log.txt")
|
||||
|
||||
# make sure we get all necessary inputs
|
||||
print("init SOURCES", cfg.data.sources)
|
||||
cfg.data.sources = expand_source_paths(cfg.data.sources)
|
||||
print("SOURCES", cfg.data.sources)
|
||||
|
||||
dataset = get_dataset_from_cfg(cfg)
|
||||
save_track_info(dataset, out_dir)
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = str(cfg.get("gpu"))
|
||||
print("CUDA_VISIBLE_DEVICES", os.environ["CUDA_VISIBLE_DEVICES"])
|
||||
device_id = cfg.get("gpu")
|
||||
|
||||
if cfg.run_opt:
|
||||
device = get_device(device_id)
|
||||
run_opt(cfg, dataset, out_dir, device)
|
||||
|
||||
if cfg.run_vis:
|
||||
run_vis(
|
||||
cfg, dataset, out_dir, device_id, **cfg.get("vis", dict())
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+475
@@ -0,0 +1,475 @@
|
||||
import os
|
||||
import glob
|
||||
|
||||
import imageio
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch.utils.data import DataLoader
|
||||
from body_model import MANO
|
||||
from omegaconf import DictConfig, OmegaConf
|
||||
|
||||
from data import get_dataset_from_cfg, expand_source_paths
|
||||
from optim.output import (
|
||||
get_results_paths,
|
||||
load_result,
|
||||
save_input_frames,
|
||||
save_input_poses,
|
||||
)
|
||||
from util.loaders import load_config_from_log, resolve_cfg_paths
|
||||
from util.tensor import get_device, move_to, detach_all, to_torch
|
||||
from vis.output import prep_result_vis, animate_scene, make_video_grid_2x2
|
||||
from vis.tools import vis_keypoints
|
||||
from vis.viewer import init_viewer
|
||||
import time
|
||||
from geometry.mesh import save_mesh_scenes, vertices_to_trimesh
|
||||
LIGHT_BLUE=(0.65098039, 0.74117647, 0.85882353)
|
||||
|
||||
def save_meshes_all(cfg, dataset, res_dicts, dev_id, mesh_dirs, num_steps=-1):
|
||||
B = len(dataset)
|
||||
T = dataset.seq_len
|
||||
loader = DataLoader(dataset, batch_size=B, shuffle=False)
|
||||
device = get_device(dev_id)
|
||||
obs_data = move_to(next(iter(loader)), device)
|
||||
|
||||
# load models
|
||||
cfg = resolve_cfg_paths(cfg)
|
||||
# Instantiate MANO model
|
||||
mano_cfg = {k.lower(): v for k,v in dict(cfg.MANO).items()}
|
||||
print('initializing MANO model with cfgs:', mano_cfg, 'B*T', B*T)
|
||||
hand_model = MANO(batch_size=B*T, pose2rot=True, **mano_cfg).to(device)
|
||||
|
||||
for res_dict, mesh_dir in zip(res_dicts, mesh_dirs):
|
||||
res_dict = move_to(res_dict, device)
|
||||
scene_dict = move_to(
|
||||
prep_result_vis(
|
||||
res_dict,
|
||||
obs_data["vis_mask"],
|
||||
obs_data["track_id"],
|
||||
hand_model,
|
||||
temporal_smooth=cfg.temporal_smooth,
|
||||
smooth_trans=True # For mesh export, smooth everything
|
||||
),
|
||||
"cpu",
|
||||
)
|
||||
|
||||
scene_dir = mesh_dir
|
||||
verts, joints, colors, l_faces, r_faces, is_right, bounds = scene_dict["geometry"]
|
||||
T = len(verts)
|
||||
print(f"{T} mesh frames for ", mesh_dir)
|
||||
times = list(range(0, T, 1))
|
||||
flag = False
|
||||
for t in times:
|
||||
if len(is_right[t]) > 1:
|
||||
flag = True
|
||||
vv = t
|
||||
|
||||
if flag:
|
||||
init_trans = (joints[vv][0][9].clone() + joints[vv][1][9].clone()) / 2
|
||||
else:
|
||||
init_trans = joints[0][0][9].clone()
|
||||
|
||||
for t in times:
|
||||
if len(is_right[t]) > 1:
|
||||
assert (is_right[t].cpu().numpy().tolist() == [0,1])
|
||||
# l_meshes = make_batch_mesh(verts[t][0][None], l_faces[t], colors[t][0][None])
|
||||
# r_meshes = make_batch_mesh(verts[t][1][None], r_faces[t], colors[t][1][None])
|
||||
# assert len(l_meshes) == 1
|
||||
# assert len(r_meshes) == 1
|
||||
|
||||
verts[t][0] -= init_trans
|
||||
joints[t][0] -= init_trans
|
||||
tmesh = vertices_to_trimesh(verts[t][0].detach().cpu().numpy(), l_faces[t].detach().cpu().numpy(), LIGHT_BLUE, is_right=0)
|
||||
tmesh.export(os.path.join(scene_dir, f'{str(t).zfill(6)}_0.obj'))
|
||||
# np.save(f'{str(t).zfill(6)}_0.npy', joints[t][0].detach().cpu().numpy())
|
||||
|
||||
verts[t][1] -= init_trans
|
||||
joints[t][1] -= init_trans
|
||||
tmesh = vertices_to_trimesh(verts[t][1].detach().cpu().numpy(), r_faces[t].detach().cpu().numpy(), LIGHT_BLUE, is_right=1)
|
||||
tmesh.export(os.path.join(scene_dir, f'{str(t).zfill(6)}_1.obj'))
|
||||
# np.save(f'{str(t).zfill(6)}_1.npy', joints[t][1].detach().cpu().numpy())
|
||||
|
||||
else:
|
||||
assert len(is_right[t]) == 1
|
||||
if is_right[t] == 0:
|
||||
verts[t][0] -= init_trans
|
||||
joints[t][0] -= init_trans
|
||||
tmesh = vertices_to_trimesh(verts[t][0].detach().cpu().numpy(), l_faces[t].detach().cpu().numpy(), LIGHT_BLUE, is_right=0)
|
||||
tmesh.export(os.path.join(scene_dir, f'{str(t).zfill(6)}_0.obj'))
|
||||
# np.save(f'{str(t).zfill(6)}_0.npy', joints[t][0].detach().cpu().numpy())
|
||||
|
||||
elif is_right[t] == 1:
|
||||
verts[t][0] -= init_trans
|
||||
joints[t][0] -= init_trans
|
||||
tmesh = vertices_to_trimesh(verts[t][0].detach().cpu().numpy(), r_faces[t].detach().cpu().numpy(), LIGHT_BLUE, is_right=1)
|
||||
tmesh.export(os.path.join(scene_dir, f'{str(t).zfill(6)}_1.obj'))
|
||||
# np.save(f'{str(t).zfill(6)}_1.npy', joints[t][0].detach().cpu().numpy())
|
||||
|
||||
|
||||
def run_vis(
|
||||
cfg,
|
||||
dataset,
|
||||
out_dir,
|
||||
dev_id,
|
||||
phases=["smooth_fit"],
|
||||
render_views=["src_cam", "above", "side"],
|
||||
make_grid=True,
|
||||
overwrite=False,
|
||||
save_dir=None,
|
||||
render_kps=False,
|
||||
render_layers=False,
|
||||
save_frames=False,
|
||||
**kwargs
|
||||
):
|
||||
save_dir = out_dir if save_dir is None else save_dir
|
||||
print("OUT_DIR", out_dir)
|
||||
print("SAVE_DIR", save_dir)
|
||||
print("VISUALIZING PHASES", phases)
|
||||
print("RENDERING VIEWS", render_views)
|
||||
print("RENDER_KPS", render_kps)
|
||||
print("OVERWRITE", overwrite)
|
||||
|
||||
# save input frames
|
||||
inp_vid_path = save_input_frames(
|
||||
dataset,
|
||||
f"{save_dir}/{dataset.seq_name}_input.mp4",
|
||||
fps=cfg.fps,
|
||||
overwrite=True,
|
||||
)
|
||||
|
||||
if render_kps:
|
||||
render_keypoints_2d(dataset, save_dir, overwrite=overwrite)
|
||||
|
||||
if len(render_views) < 1:
|
||||
return
|
||||
|
||||
out_ext = "/" if render_layers or save_frames else ".mp4"
|
||||
phase_results = {}
|
||||
phase_max_iters = {}
|
||||
for phase in phases:
|
||||
res_dir = os.path.join(out_dir, phase)
|
||||
if phase == "input":
|
||||
res = get_input_dict(dataset)
|
||||
it = f"{0:06d}"
|
||||
|
||||
elif os.path.isdir(res_dir):
|
||||
res_path_dict = get_results_paths(res_dir)
|
||||
print(f"FOUND {len(res_path_dict)} results in {res_dir}")
|
||||
it = sorted(res_path_dict.keys())[-1]
|
||||
res = load_result(res_path_dict[it])["world"]
|
||||
|
||||
else:
|
||||
print(f"{res_dir} does not exist, skipping")
|
||||
continue
|
||||
|
||||
out_name = f"{save_dir}/{dataset.seq_name}_{phase}_final_{it}"
|
||||
mesh_dir = f"{save_dir}/{phase}/{dataset.seq_name}_{it}_meshes"
|
||||
os.makedirs(mesh_dir, exist_ok=True)
|
||||
phase_max_iters[phase] = it
|
||||
|
||||
out_paths = [f"{out_name}_{view}{out_ext}" for view in render_views]
|
||||
if not overwrite and all(os.path.exists(p) for p in out_paths):
|
||||
print("FOUND OUT PATHS", out_paths)
|
||||
continue
|
||||
|
||||
phase_results[phase] = out_name, mesh_dir, res
|
||||
|
||||
if len(phase_results) > 0:
|
||||
out_names, mesh_dir, res_dicts = zip(*phase_results.values())
|
||||
render_results(
|
||||
cfg,
|
||||
dataset,
|
||||
dev_id,
|
||||
res_dicts,
|
||||
out_names,
|
||||
render_views=render_views,
|
||||
render_layers=render_layers,
|
||||
save_frames=save_frames,
|
||||
**kwargs,
|
||||
)
|
||||
save_meshes_all(cfg, dataset, res_dicts, dev_id, mesh_dir)
|
||||
|
||||
if make_grid:
|
||||
for phase, it in phase_max_iters.items():
|
||||
grid_path = f"{save_dir}/{dataset.seq_name}_{phase}_grid.mp4"
|
||||
vid_paths = [
|
||||
f"{save_dir}/{dataset.seq_name}_{phase}_final_{it}_src_cam.mp4",
|
||||
f"{save_dir}/{dataset.seq_name}_{phase}_final_{it}_front.mp4",
|
||||
f"{save_dir}/{dataset.seq_name}_{phase}_final_{it}_above.mp4",
|
||||
f"{save_dir}/{dataset.seq_name}_{phase}_final_{it}_side.mp4",
|
||||
]
|
||||
make_video_grid_2x2(
|
||||
grid_path,
|
||||
vid_paths,
|
||||
fps=cfg.fps,
|
||||
overwrite=True,
|
||||
)
|
||||
|
||||
|
||||
def get_input_dict(dataset):
|
||||
dataset.load_data(interp_input=False)
|
||||
d = dataset.data_dict
|
||||
input_params = {
|
||||
"pose_body": np.stack(d["init_body_pose"], axis=0),
|
||||
"trans": np.stack(d["init_trans"], axis=0),
|
||||
"root_orient": np.stack(d["init_root_orient"], axis=0),
|
||||
}
|
||||
input_params = to_torch(input_params)
|
||||
print({k: v.shape for k, v in input_params.items()})
|
||||
return input_params
|
||||
|
||||
|
||||
def render_keypoints_2d(dataset, save_dir, overwrite=False):
|
||||
"""
|
||||
render 2d keypoints for each track
|
||||
"""
|
||||
dataset.load_data()
|
||||
out_dir = f"{save_dir}/{dataset.seq_name}_joints2d"
|
||||
B, T = dataset.n_tracks, dataset.seq_len
|
||||
if not overwrite and (os.path.isdir(out_dir) and len(os.listdir(out_dir)) >= B * T):
|
||||
print(f"Keypoints already rendered in {out_dir}")
|
||||
return
|
||||
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
for i, tid in enumerate(dataset.track_ids):
|
||||
joints2d = dataset.data_dict["joints2d"][i] # (T, J, 3)
|
||||
for t, sel_img_name in enumerate(dataset.sel_img_names):
|
||||
img = vis_keypoints(joints2d[t : t + 1], dataset.img_size)
|
||||
out_path = f"{out_dir}/{sel_img_name}_{tid}.png"
|
||||
imageio.imwrite(out_path, img)
|
||||
|
||||
|
||||
|
||||
def render_results(cfg, dataset, dev_id, res_dicts, out_names, **kwargs):
|
||||
"""
|
||||
render results for all selected phases
|
||||
"""
|
||||
assert len(res_dicts) == len(out_names)
|
||||
if len(res_dicts) < 1:
|
||||
print("no results to render, skipping")
|
||||
return
|
||||
|
||||
B = len(dataset)
|
||||
T = dataset.seq_len
|
||||
loader = DataLoader(dataset, batch_size=B, shuffle=False)
|
||||
|
||||
device = get_device(dev_id)
|
||||
obs_data = move_to(next(iter(loader)), device)
|
||||
cam_data = dataset.get_camera_data()
|
||||
|
||||
# load models
|
||||
cfg = resolve_cfg_paths(cfg)
|
||||
# Instantiate MANO model
|
||||
mano_cfg = {k.lower(): v for k,v in dict(cfg.MANO).items()}
|
||||
print('initializing MANO model with cfgs:', mano_cfg, 'B*T', B*T)
|
||||
hand_model = MANO(batch_size=B*T, pose2rot=True, **mano_cfg).to(device)
|
||||
vis = init_viewer(
|
||||
dataset.img_size,
|
||||
cam_data["intrins"][0],
|
||||
vis_scale=1.0,
|
||||
bg_paths=dataset.sel_img_paths,
|
||||
fps=cfg.fps,
|
||||
)
|
||||
|
||||
# Set 2D keypoints for overlay visualization if render_keypoints is enabled
|
||||
if kwargs.get('render_keypoints', False):
|
||||
# Load 2D keypoints from dataset - assume single track for now
|
||||
# joints2d is list of (T, J, 3) arrays, one per track
|
||||
dataset.load_data()
|
||||
if len(dataset.data_dict["joints2d"]) > 0:
|
||||
joints2d = dataset.data_dict["joints2d"][0] # First track: (T, J, 3)
|
||||
keypoints_seq = [joints2d[t] for t in range(T)] # List of (J, 3) arrays
|
||||
vis.set_keypoints_seq(keypoints_seq)
|
||||
print(f"Loaded {T} frames of 2D keypoints with {joints2d.shape[1]} joints")
|
||||
|
||||
save_paths_all = []
|
||||
render_views = kwargs.get('render_views', ['src_cam', 'above', 'side'])
|
||||
|
||||
# Separate src_cam from other views
|
||||
src_cam_views = [v for v in render_views if v == 'src_cam']
|
||||
other_views = [v for v in render_views if v != 'src_cam']
|
||||
|
||||
for res_dict, out_name in zip(res_dicts, out_names):
|
||||
print(f'preparing results for rendering {out_name}')
|
||||
res_dict = move_to(res_dict, device)
|
||||
|
||||
# Render src_cam WITH temporal smoothing but WITHOUT trans smoothing
|
||||
if src_cam_views:
|
||||
print(f"Rendering src_cam view WITH temporal smoothing (excluding trans)")
|
||||
scene_dict_no_smooth = prep_result_vis(
|
||||
res_dict,
|
||||
obs_data["vis_mask"],
|
||||
obs_data["track_id"],
|
||||
hand_model,
|
||||
temporal_smooth=cfg.temporal_smooth, # Apply smoothing
|
||||
smooth_trans=False # But don't smooth translation
|
||||
)
|
||||
|
||||
# Compute predicted 2D keypoints using THE SAME function as optimization
|
||||
if kwargs.get('render_keypoints', False):
|
||||
from geometry import camera as cam_util
|
||||
from body_model import run_mano
|
||||
|
||||
# Get the same data as optimization
|
||||
joints3d_op = run_mano(
|
||||
hand_model,
|
||||
res_dict["trans"],
|
||||
res_dict["root_orient"],
|
||||
res_dict["pose_body"],
|
||||
res_dict["is_right"],
|
||||
res_dict.get("betas", None),
|
||||
)["joints"] # (B, T, J, 3)
|
||||
|
||||
# Get cameras (same as optimization)
|
||||
cam_R = res_dict["cam_R"] # (B, T, 3, 3)
|
||||
cam_t = res_dict["cam_t"] # (B, T, 3)
|
||||
intrins = res_dict["intrins"] # (4,) or (T, 4) or (B, T, 4)
|
||||
|
||||
# Ensure intrins is (B, T, 4)
|
||||
if intrins.ndim == 1: # (4,)
|
||||
intrins = intrins[None, None].expand(cam_R.shape[0], cam_R.shape[1], -1) # (B, T, 4)
|
||||
elif intrins.ndim == 2: # (T, 4)
|
||||
intrins = intrins[None].expand(cam_R.shape[0], -1, -1) # (B, T, 4)
|
||||
|
||||
cam_f = intrins[:, :, :2] # (B, T, 2)
|
||||
cam_center = intrins[:, :, 2:] # (B, T, 2)
|
||||
|
||||
# Reproject using EXACTLY the same function as optimization
|
||||
joints2d_pred = cam_util.reproject(joints3d_op, cam_R, cam_t, cam_f[0], cam_center[0]) # (B, T, J, 2)
|
||||
|
||||
# Convert to list for first track
|
||||
pred_keypoints_seq = [joints2d_pred[0, t].cpu().numpy() for t in range(joints2d_pred.shape[1])]
|
||||
vis.set_pred_keypoints_seq(pred_keypoints_seq)
|
||||
print(f"Computed {len(pred_keypoints_seq)} frames of predicted 2D keypoints using reproject()")
|
||||
|
||||
# Render only src_cam view
|
||||
kwargs_src = kwargs.copy()
|
||||
kwargs_src['render_views'] = src_cam_views
|
||||
save_paths_src = animate_scene(
|
||||
vis, scene_dict_no_smooth, out_name, seq_name=dataset.seq_name, **kwargs_src
|
||||
)
|
||||
save_paths_all.append(save_paths_src)
|
||||
|
||||
# Render other views WITH temporal smoothing (including trans)
|
||||
if other_views:
|
||||
print(f"Rendering {other_views} views WITH temporal smoothing (including trans)")
|
||||
scene_dict_smooth = prep_result_vis(
|
||||
res_dict,
|
||||
obs_data["vis_mask"],
|
||||
obs_data["track_id"],
|
||||
hand_model,
|
||||
temporal_smooth=cfg.temporal_smooth, # Apply smoothing for other views
|
||||
smooth_trans=True # Smooth translation for other views
|
||||
)
|
||||
|
||||
# Render other views
|
||||
kwargs_other = kwargs.copy()
|
||||
kwargs_other['render_views'] = other_views
|
||||
kwargs_other['render_keypoints'] = False # No keypoints on other views
|
||||
save_paths_other = animate_scene(
|
||||
vis, scene_dict_smooth, out_name, seq_name=dataset.seq_name, **kwargs_other
|
||||
)
|
||||
if not src_cam_views:
|
||||
save_paths_all.append(save_paths_other)
|
||||
|
||||
vis.close()
|
||||
|
||||
return save_paths_all
|
||||
|
||||
|
||||
def visualize_log(log_dir, dev_id, phases, save_dir=None, **kwargs):
|
||||
print(log_dir)
|
||||
cfg = load_config_from_log(log_dir)
|
||||
print(cfg)
|
||||
print(cfg.data)
|
||||
print(cfg.data.sources)
|
||||
|
||||
# make sure we get all necessary inputs
|
||||
cfg.data.sources = expand_source_paths(cfg.data.sources)
|
||||
print("SOURCES", cfg.data.sources)
|
||||
# cfg.data.track_ids = "001"
|
||||
dataset = get_dataset_from_cfg(cfg)
|
||||
if len(dataset) < 1:
|
||||
print(f"No tracks in dataset, skipping")
|
||||
return
|
||||
|
||||
run_vis(cfg, dataset, log_dir, dev_id, phases=phases, save_dir=save_dir, **kwargs)
|
||||
|
||||
|
||||
def launch_vis(i, args):
|
||||
log_dir = args.log_dirs[i]
|
||||
dev_id = args.gpus[i % len(args.gpus)]
|
||||
os.environ["EGL_DEVICE_ID"] = str(dev_id)
|
||||
os.environ["PYOPENGL_PLATFORM"] = "egl"
|
||||
|
||||
if args.save_root is not None:
|
||||
path_name = log_dir.split(args.log_root)[-1].strip("/")
|
||||
exp_name = "-".join(path_name.split("/")[:2])
|
||||
save_dir = f"{args.save_root}/{exp_name}"
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
|
||||
print('args.save_root: ', args.save_root)
|
||||
visualize_log(
|
||||
log_dir,
|
||||
dev_id,
|
||||
phases=args.phases,
|
||||
save_dir=save_dir,
|
||||
overwrite=args.overwrite,
|
||||
accumulate=args.accumulate,
|
||||
render_kps=args.render_kps,
|
||||
render_layers=args.render_layers,
|
||||
render_views=args.render_views,
|
||||
save_frames=args.save_frames,
|
||||
make_grid=args.grid,
|
||||
)
|
||||
|
||||
|
||||
def main(args):
|
||||
"""
|
||||
visualize all runs in root
|
||||
"""
|
||||
OmegaConf.register_new_resolver("eval", eval)
|
||||
log_dirs = []
|
||||
for root, subd, files in os.walk(args.log_root):
|
||||
if ".hydra" in subd:
|
||||
log_dirs.append(root)
|
||||
args.log_dirs = log_dirs
|
||||
print(f"FOUND {len(args.log_dirs)} TO RENDER")
|
||||
|
||||
if len(args.gpus) > 1:
|
||||
from torch.multiprocessing import Pool
|
||||
|
||||
torch.multiprocessing.set_start_method("spawn")
|
||||
|
||||
with Pool(processes=len(args.gpus)) as pool:
|
||||
res = pool.starmap(
|
||||
launch_vis, [(i, args) for i in range(len(args.log_dirs))]
|
||||
)
|
||||
return
|
||||
|
||||
for i in range(len(args.log_dirs)):
|
||||
launch_vis(i, args)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--log_root", required=True)
|
||||
parser.add_argument("--save_root", default='./')
|
||||
parser.add_argument("--phases", nargs="*", default=["smooth_fit"])
|
||||
parser.add_argument("--gpus", nargs="*", default=[0])
|
||||
parser.add_argument(
|
||||
"-rv",
|
||||
"--render_views",
|
||||
nargs="*",
|
||||
default=["src_cam", "front", "above", "side"],
|
||||
)
|
||||
parser.add_argument("-g", "--grid", action="store_true")
|
||||
parser.add_argument("-rl", "--render_layers", action="store_true")
|
||||
parser.add_argument("-kp", "--render_kps", action="store_true")
|
||||
parser.add_argument("-sf", "--save_frames", action="store_true")
|
||||
parser.add_argument("-ra", "--accumulate", action="store_true")
|
||||
parser.add_argument("-y", "--overwrite", action="store_true")
|
||||
args = parser.parse_args()
|
||||
|
||||
main(args)
|
||||
+54
@@ -0,0 +1,54 @@
|
||||
#!/bin/bash
|
||||
|
||||
# Point BASE_PATH directly to subject1
|
||||
BASE_PATH="/data/home/zy3023/code/hand/slam-hand/test/HOT3D"
|
||||
VIDEO_JSON="/data/home/zy3023/code/hand/slam-hand/slahmr-eccv/confs/data/video.yaml"
|
||||
|
||||
# Directory for log files
|
||||
LOG_DIR="logs"
|
||||
mkdir -p "$LOG_DIR"
|
||||
|
||||
# How many Python processes to run in parallel at most
|
||||
MAX_JOBS=2
|
||||
|
||||
count=0 # Tracks how many jobs have been launched
|
||||
|
||||
for mp4file in "$BASE_PATH"/*.mp4; do
|
||||
[ -f "$mp4file" ] || continue
|
||||
|
||||
# basename "a.mp4" => "a.mp4"
|
||||
mp4basename=$(basename "$mp4file")
|
||||
# strip the ".mp4" extension => "a"
|
||||
vidname="${mp4basename%.mp4}"
|
||||
|
||||
logfile="${LOG_DIR}/${vidname}.log"
|
||||
|
||||
# Skip this video if the log file exists and is non-empty
|
||||
if [ -s "$logfile" ]; then
|
||||
echo "Skipping $vidname, log file exists and is non-empty: $logfile"
|
||||
continue
|
||||
fi
|
||||
|
||||
# Update 'seq' in the YAML to "a" instead of "a.mp4"
|
||||
sed -i "s|seq: *\"[^\"]*\"|seq: \"$vidname\"|" "$VIDEO_JSON"
|
||||
|
||||
echo "Launching job for: $vidname => seq: \"$vidname\". Logs: $logfile"
|
||||
|
||||
# Run Python in the background (&).
|
||||
# Output goes to "$logfile"
|
||||
python -u run_opt.py data=video run_opt=True run_vis=True data.seq=$vidname \
|
||||
> "$logfile" 2>&1 &
|
||||
|
||||
((count++))
|
||||
# If we've launched MAX_JOBS in parallel, wait for them to complete
|
||||
if (( count % MAX_JOBS == 0 )); then
|
||||
echo "Reached $count jobs. Waiting for them to finish..."
|
||||
wait
|
||||
echo "Continuing..."
|
||||
fi
|
||||
done
|
||||
|
||||
# After the loop, wait for any remaining jobs
|
||||
wait
|
||||
|
||||
echo "All done."
|
||||
@@ -0,0 +1,3 @@
|
||||
from . import loaders
|
||||
from . import logger
|
||||
from . import tensor
|
||||
+207
@@ -0,0 +1,207 @@
|
||||
"""
|
||||
Utilities for loading input and models (reading pre-processed data)
|
||||
"""
|
||||
|
||||
import os
|
||||
import importlib
|
||||
import glob
|
||||
from omegaconf import OmegaConf
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
# from body_model import HandModel
|
||||
|
||||
|
||||
ROOT_DIR = os.path.abspath(f"{__file__}/../../../")
|
||||
|
||||
|
||||
def load_config_from_log(log_dir):
|
||||
hydra_dir = f"{log_dir}/.hydra"
|
||||
cfg_path = f"{hydra_dir}/config.yaml"
|
||||
assert os.path.isdir(hydra_dir), f"{hydra_dir} does not exist"
|
||||
assert os.path.isfile(cfg_path), f"{cfg_path} does not exist"
|
||||
return OmegaConf.load(cfg_path)
|
||||
|
||||
|
||||
def resolve_cfg_paths(cfg):
|
||||
paths = cfg.paths
|
||||
for name, rel_path in paths.items():
|
||||
if rel_path.startswith("/"): # absolute path
|
||||
continue
|
||||
paths[name] = os.path.join(ROOT_DIR, rel_path)
|
||||
print("RESOLVED PATHS", paths)
|
||||
return cfg
|
||||
|
||||
|
||||
# def load_mano_body_model(cfg):
|
||||
# """
|
||||
# Load SMPL model
|
||||
# """
|
||||
# if device is None:
|
||||
# device = torch.device("cpu")
|
||||
# return (
|
||||
# HandModel(
|
||||
# bm_path=path,
|
||||
# num_betas=num_betas,
|
||||
# batch_size=batch_size,
|
||||
# use_vtx_selector=use_vtx_selector,
|
||||
# model_type=model_type,
|
||||
# ).to(device),
|
||||
# )
|
||||
|
||||
|
||||
def expid2model(expr_dir):
|
||||
""" "
|
||||
Reading VPoser models (https://github.com/nghorbani/human_body_prior).
|
||||
"""
|
||||
from configer import Configer
|
||||
|
||||
if not os.path.exists(expr_dir):
|
||||
raise ValueError("Could not find the experiment directory: %s" % expr_dir)
|
||||
|
||||
best_model_fname = sorted(
|
||||
glob.glob(os.path.join(expr_dir, "snapshots", "*.pt")), key=os.path.getmtime
|
||||
)[-1]
|
||||
try_num = os.path.basename(best_model_fname).split("_")[0]
|
||||
|
||||
print(("Found Trained Model: %s" % best_model_fname))
|
||||
|
||||
default_ps_fname = glob.glob(os.path.join(expr_dir, "*.ini"))[0]
|
||||
if not os.path.exists(default_ps_fname):
|
||||
raise ValueError(
|
||||
"Could not find the appropriate vposer_settings: %s" % default_ps_fname
|
||||
)
|
||||
ps = Configer(
|
||||
default_ps_fname=default_ps_fname,
|
||||
work_dir=expr_dir,
|
||||
best_model_fname=best_model_fname,
|
||||
)
|
||||
|
||||
return ps, best_model_fname
|
||||
|
||||
|
||||
def load_vposer(expr_dir, vp_model="snapshot"):
|
||||
"""
|
||||
:param expr_dir:
|
||||
:param vp_model: either 'snapshot' to use the experiment folder's code or a VPoser imported module, e.g.
|
||||
from human_body_prior.train.vposer_smpl import VPoser, then pass VPoser to this function
|
||||
:param if True will load the model definition used for training, and not the one in current repository
|
||||
:return:
|
||||
"""
|
||||
ps, trained_model_fname = expid2model(expr_dir)
|
||||
if vp_model == "snapshot":
|
||||
|
||||
vposer_path = sorted(
|
||||
glob.glob(os.path.join(expr_dir, "vposer_*.py")), key=os.path.getmtime
|
||||
)[-1]
|
||||
|
||||
spec = importlib.util.spec_from_file_location("VPoser", vposer_path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(module)
|
||||
|
||||
vposer_pt = getattr(module, "VPoser")(
|
||||
num_neurons=ps.num_neurons, latentD=ps.latentD, data_shape=ps.data_shape
|
||||
)
|
||||
else:
|
||||
vposer_pt = vp_model(
|
||||
num_neurons=ps.num_neurons, latentD=ps.latentD, data_shape=ps.data_shape
|
||||
)
|
||||
|
||||
vposer_pt.load_state_dict(torch.load(trained_model_fname, map_location="cpu"))
|
||||
vposer_pt.eval()
|
||||
|
||||
return vposer_pt, ps
|
||||
|
||||
|
||||
def load_state(
|
||||
load_path,
|
||||
model,
|
||||
optimizer=None,
|
||||
is_parallel=False,
|
||||
map_location=None,
|
||||
ignore_keys=None,
|
||||
):
|
||||
"""
|
||||
Load Humor model checkpoint
|
||||
"""
|
||||
if not os.path.exists(load_path):
|
||||
print("Could not find checkpoint at path " + load_path)
|
||||
|
||||
full_checkpoint_dict = torch.load(load_path, map_location=map_location)
|
||||
model_state_dict = full_checkpoint_dict["model"]
|
||||
optim_state_dict = full_checkpoint_dict["optim"]
|
||||
|
||||
# load model weights
|
||||
for k, v in model_state_dict.items():
|
||||
if k.split(".")[0] == "module" and not is_parallel:
|
||||
# then it was trained with Data parallel
|
||||
print("Loading weights trained with DataParallel...")
|
||||
model_state_dict = {
|
||||
".".join(k.split(".")[1:]): v
|
||||
for k, v in model_state_dict.items()
|
||||
if k.split(".")[0] == "module"
|
||||
}
|
||||
break
|
||||
|
||||
if ignore_keys is not None:
|
||||
model_state_dict = {
|
||||
k: v
|
||||
for k, v in model_state_dict.items()
|
||||
if k.split(".")[0] not in ignore_keys
|
||||
}
|
||||
|
||||
# overwrite entries in the existing state dict
|
||||
missing_keys, unexpected_keys = model.load_state_dict(
|
||||
model_state_dict, strict=False
|
||||
)
|
||||
if ignore_keys is not None:
|
||||
missing_keys = [k for k in missing_keys if k.split(".")[0] not in ignore_keys]
|
||||
unexpected_keys = [
|
||||
k for k in unexpected_keys if k.split(".")[0] not in ignore_keys
|
||||
]
|
||||
if len(missing_keys) > 0:
|
||||
print(
|
||||
"WARNING: The following keys could not be found in the given state dict - ignoring..."
|
||||
)
|
||||
print(missing_keys)
|
||||
if len(unexpected_keys) > 0:
|
||||
print(
|
||||
"WARNING: The following keys were found in the given state dict but not in the current model - ignoring..."
|
||||
)
|
||||
print(unexpected_keys)
|
||||
|
||||
# load optimizer weights
|
||||
if optimizer is not None:
|
||||
optimizer.load_state_dict(optim_state_dict)
|
||||
|
||||
min_train_loss = float("Inf")
|
||||
if "min_train_loss" in full_checkpoint_dict.keys():
|
||||
min_train_loss = full_checkpoint_dict["min_train_loss"]
|
||||
|
||||
return (
|
||||
full_checkpoint_dict["epoch"],
|
||||
full_checkpoint_dict["min_val_loss"],
|
||||
min_train_loss,
|
||||
)
|
||||
|
||||
|
||||
def load_gmm(path, device=None):
|
||||
"""
|
||||
Load motion prior GMM for first frame motion prior
|
||||
"""
|
||||
if device is None:
|
||||
device = torch.device("cpu")
|
||||
|
||||
gmm_path = os.path.join(path, "prior_gmm.npz")
|
||||
if not os.path.isfile(gmm_path):
|
||||
raise ValueError(f"Init motion prior path {gmm_path} does not exist")
|
||||
|
||||
res = np.load(gmm_path)
|
||||
return {
|
||||
"gmm": (
|
||||
torch.from_numpy(res["weights"]).to(device),
|
||||
torch.from_numpy(res["means"]).to(device),
|
||||
torch.from_numpy(res["covariances"]).to(device),
|
||||
)
|
||||
}
|
||||
+37
@@ -0,0 +1,37 @@
|
||||
import datetime
|
||||
|
||||
|
||||
class Logger(object):
|
||||
"""
|
||||
Static logging class
|
||||
"""
|
||||
|
||||
log_file = None
|
||||
|
||||
@staticmethod
|
||||
def init(log_path):
|
||||
Logger.log_file = log_path
|
||||
|
||||
@staticmethod
|
||||
def log(write_str, to_stdout=True):
|
||||
if to_stdout:
|
||||
print(write_str)
|
||||
if not Logger.log_file:
|
||||
print("Logger must be initialized before logging!")
|
||||
return
|
||||
time_str = datetime.datetime.now().strftime("%Y-%m-%d_%H:%M:%S")
|
||||
with open(Logger.log_file, "a") as f:
|
||||
f.write(time_str + " ")
|
||||
f.write(str(write_str) + "\n")
|
||||
|
||||
|
||||
def log_cur_stats(stats_dict, iter=None, to_stdout=True):
|
||||
loss = stats_dict.pop("total", 0)
|
||||
Logger.log(f"LOSS: {loss:.04f}", to_stdout=to_stdout)
|
||||
for k, v in stats_dict.items():
|
||||
Logger.log(f"{k}: {v:.04f}", to_stdout=to_stdout)
|
||||
if to_stdout:
|
||||
if iter is not None:
|
||||
print("======= iter %d =======" % iter)
|
||||
else:
|
||||
print("========")
|
||||
+139
@@ -0,0 +1,139 @@
|
||||
import torch
|
||||
import numpy as np
|
||||
|
||||
|
||||
def move_to(obj, device):
|
||||
if isinstance(obj, torch.Tensor):
|
||||
return obj.to(device)
|
||||
if isinstance(obj, dict):
|
||||
return {k: move_to(v, device) for k, v in obj.items()}
|
||||
if isinstance(obj, (list, tuple)):
|
||||
return [move_to(x, device) for x in obj]
|
||||
return obj # otherwise do nothing
|
||||
|
||||
|
||||
def detach_all(obj):
|
||||
if isinstance(obj, torch.Tensor):
|
||||
return obj.detach()
|
||||
if isinstance(obj, dict):
|
||||
return {k: detach_all(v) for k, v in obj.items()}
|
||||
if isinstance(obj, (list, tuple)):
|
||||
return [detach_all(x) for x in obj]
|
||||
return obj # otherwise do nothing
|
||||
|
||||
|
||||
def to_torch(obj):
|
||||
if isinstance(obj, np.ndarray):
|
||||
return torch.from_numpy(obj).float()
|
||||
if isinstance(obj, dict):
|
||||
return {k: to_torch(v) for k, v in obj.items()}
|
||||
if isinstance(obj, (list, tuple)):
|
||||
return [to_torch(x) for x in obj]
|
||||
return obj
|
||||
|
||||
|
||||
def to_np(obj):
|
||||
if isinstance(obj, torch.Tensor):
|
||||
return obj.numpy()
|
||||
if isinstance(obj, dict):
|
||||
return {k: to_np(v) for k, v in obj.items()}
|
||||
if isinstance(obj, (list, tuple)):
|
||||
return [to_np(x) for x in obj]
|
||||
return obj
|
||||
|
||||
|
||||
def get_device(i=0):
|
||||
device = f"cuda:{i}" if torch.cuda.is_available() else "cpu"
|
||||
return torch.device(device)
|
||||
|
||||
|
||||
def scatter_intervals(tensor, start, end, T=None):
|
||||
"""
|
||||
Scatter the tensor contents into intervals from start to end
|
||||
output tensor indexed from 0 to end.max()
|
||||
:param tensor (B, S, *)
|
||||
:param start (B) start indices
|
||||
:param end (B) end indices
|
||||
:param T (int, optional) max length
|
||||
returns (B, T, *) scattered tensor
|
||||
"""
|
||||
assert isinstance(tensor, torch.Tensor) and tensor.ndim >= 2
|
||||
if T is not None:
|
||||
end = torch.where(end < 0, T + 1 + end, end)
|
||||
assert torch.all(end <= T)
|
||||
else:
|
||||
T = end.max()
|
||||
|
||||
B, S, *dims = tensor.shape
|
||||
start, end = start.long(), end.long()
|
||||
# get idcs that go past the last time step so we don't have repeat indices in scatter
|
||||
idcs = time_segment_idcs(start, end, clip=False) # (B, S)
|
||||
# mask out the extra padding
|
||||
mask = idcs >= end[:, None]
|
||||
tensor[mask] = 0
|
||||
|
||||
idcs = idcs.reshape(B, S, *(1,) * len(dims)).repeat(1, 1, *dims)
|
||||
output = torch.zeros(
|
||||
B, idcs.max() + 1, *dims, device=tensor.device, dtype=tensor.dtype
|
||||
)
|
||||
output.scatter_(1, idcs, tensor)
|
||||
# slice out the extra segments
|
||||
return output[:, :T]
|
||||
|
||||
|
||||
def get_scatter_mask(start, end, T):
|
||||
"""
|
||||
get the mask of selected intervals
|
||||
"""
|
||||
B = start.shape[0]
|
||||
start, end = start.long(), end.long()
|
||||
assert torch.all(end <= T)
|
||||
idcs = time_segment_idcs(start, end, clip=True)
|
||||
mask = torch.zeros(B, T, device=start.device, dtype=torch.bool)
|
||||
mask.scatter_(1, idcs, 1)
|
||||
return mask
|
||||
|
||||
|
||||
def select_intervals(series, start, end):
|
||||
"""
|
||||
Select slices of a tensor from start to end
|
||||
will pad uneven sequences to all the max segment length
|
||||
:param series (B, T, *)
|
||||
:param start (B)
|
||||
:param end (B)
|
||||
returns (B, S, *) selected segments, S = max(end - start)
|
||||
"""
|
||||
B, T, *dims = series.shape
|
||||
assert torch.all(end <= T)
|
||||
sel = time_segment_idcs(start, end, clip=True)
|
||||
S = sel.shape[1]
|
||||
sel = sel.reshape(B, S, *(1,) * len(dims)).repeat(1, 1, *dims)
|
||||
return torch.gather(series, 1, sel)
|
||||
|
||||
|
||||
def get_select_mask(start, end):
|
||||
"""
|
||||
get the mask of unpadded elementes for the selected time segments
|
||||
e.g. sel[mask] are the unpadded elements
|
||||
:param start (B)
|
||||
:param end (B)
|
||||
"""
|
||||
idcs = time_segment_idcs(start, end, clip=False)
|
||||
return idcs < end[:, None] # (B, S)
|
||||
|
||||
|
||||
def time_segment_idcs(start, end, clip=True):
|
||||
"""
|
||||
:param start (B)
|
||||
:param end (B)
|
||||
returns (B, S) long tensor of indices, where S = max(end - start)
|
||||
"""
|
||||
start, end = start.long(), end.long()
|
||||
S = (end - start).max()
|
||||
seg = torch.arange(S, dtype=torch.int64, device=start.device)
|
||||
idcs = start[:, None] + seg[None, :] # (B, S)
|
||||
if clip:
|
||||
# clip at the lengths of each track
|
||||
imax = end[:, None] - 1
|
||||
idcs = idcs.clamp(max=imax)
|
||||
return idcs
|
||||
@@ -0,0 +1,4 @@
|
||||
from . import viewer
|
||||
|
||||
# from . import render
|
||||
from . import tools
|
||||
+98
@@ -0,0 +1,98 @@
|
||||
241 91 181
|
||||
0 187 249
|
||||
254 228 64
|
||||
0 245 212
|
||||
222 158 54
|
||||
255 133 141
|
||||
93 169 233
|
||||
155 93 229
|
||||
24 242 178
|
||||
252 116 186
|
||||
138 225 252
|
||||
178 171 242
|
||||
24 231 98
|
||||
232 93 117
|
||||
207 250 231
|
||||
253 208 242
|
||||
109 157 197
|
||||
166 189 219
|
||||
253 146 207
|
||||
167 201 87
|
||||
117 68 177
|
||||
255 229 50
|
||||
251 202 239
|
||||
58 134 255
|
||||
255 0 110
|
||||
251 86 7
|
||||
188 51 209
|
||||
122 229 130
|
||||
0 48 73
|
||||
214 40 40
|
||||
229 179 179
|
||||
0 187 249
|
||||
255 190 11
|
||||
204 213 174
|
||||
0 245 212
|
||||
255 153 200
|
||||
144 251 146
|
||||
189 211 147
|
||||
230 0 86
|
||||
0 95 57
|
||||
0 174 126
|
||||
255 116 163
|
||||
189 198 255
|
||||
90 219 255
|
||||
158 0 142
|
||||
255 147 126
|
||||
164 36 0
|
||||
0 21 68
|
||||
145 208 203
|
||||
95 173 78
|
||||
107 104 130
|
||||
0 125 181
|
||||
106 130 108
|
||||
252 246 189
|
||||
208 244 222
|
||||
169 222 249
|
||||
228 193 249
|
||||
122 204 174
|
||||
194 140 159
|
||||
0 143 156
|
||||
235 0 0
|
||||
255 2 157
|
||||
104 61 59
|
||||
150 138 232
|
||||
152 255 82
|
||||
167 87 64
|
||||
1 255 254
|
||||
255 238 232
|
||||
254 137 0
|
||||
1 208 255
|
||||
187 136 0
|
||||
165 255 210
|
||||
255 166 254
|
||||
119 77 0
|
||||
122 71 130
|
||||
38 52 0
|
||||
0 71 84
|
||||
67 0 44
|
||||
181 0 255
|
||||
255 177 103
|
||||
255 219 102
|
||||
126 45 210
|
||||
229 111 254
|
||||
222 255 116
|
||||
0 255 120
|
||||
0 155 255
|
||||
0 100 1
|
||||
0 118 255
|
||||
133 169 0
|
||||
0 185 23
|
||||
120 130 49
|
||||
0 255 198
|
||||
255 110 65
|
||||
232 94 190
|
||||
1 0 103
|
||||
149 0 58
|
||||
98 14 0
|
||||
0 0 0
|
||||
+245
@@ -0,0 +1,245 @@
|
||||
import torch
|
||||
import numpy as np
|
||||
|
||||
|
||||
def get_seq_figure_skip(seq_name=None):
|
||||
if seq_name == "017437_mpii_test":
|
||||
return 15
|
||||
if seq_name == "012968_mpii_test":
|
||||
return 15
|
||||
if seq_name == "recording_20211002_S03_S18_04-all-100-200":
|
||||
return 11
|
||||
return 10
|
||||
|
||||
|
||||
def get_seq_static_lookat_points(seq_name=None, bounds=None):
|
||||
if seq_name == "002276_mpii_test":
|
||||
top_source = torch.tensor([-0.5, -3.0, -2.0])
|
||||
top_target = torch.tensor([0.0, 0.0, 4.0])
|
||||
|
||||
side_source = torch.tensor([6.0, -2.0, 2.0])
|
||||
side_target = torch.tensor([2.0, 0.0, 6.0])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "002374_mpii_test":
|
||||
top_source = torch.tensor([-4.0, 0.0, 0.0])
|
||||
top_target = torch.tensor([-2.0, 2.0, 3.0])
|
||||
|
||||
side_source = torch.tensor([-7.0, 1.0, 7.0])
|
||||
side_target = torch.tensor([0.0, 1.0, 8.0])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "003742_mpii_test":
|
||||
top_source = torch.tensor([3.0, -2.0, -6])
|
||||
top_target = torch.tensor([3.0, 1.0, 2.0])
|
||||
|
||||
side_source = torch.tensor([-5.0, -1.0, -2])
|
||||
side_target = torch.tensor([2.0, 0.0, 5])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "003747_mpii_test":
|
||||
top_source = torch.tensor([-5.0, -2.0, -5])
|
||||
top_target = torch.tensor([-3.0, 2.0, 1.0])
|
||||
|
||||
side_source = torch.tensor([-9.0, -1.0, -4])
|
||||
side_target = torch.tensor([-3.0, 2.0, 3])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "003943_mpii_test":
|
||||
top_source = torch.tensor([0, -1, -2])
|
||||
top_target = torch.tensor([0, 1, 3])
|
||||
|
||||
side_source = torch.tensor([6, 0, 1])
|
||||
side_target = torch.tensor([0, 0, 5])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "006537_mpii_test":
|
||||
top_source = torch.tensor([0, -4, -1])
|
||||
top_target = torch.tensor([1.0, 0.0, 5.0])
|
||||
|
||||
side_source = torch.tensor([5, -2, 0])
|
||||
side_target = torch.tensor([1.0, 1.0, 7.0])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "007684_mpii_test":
|
||||
top_source = torch.tensor([-2, -4, -6])
|
||||
top_target = torch.tensor([0.5, 0.0, 4.0])
|
||||
|
||||
side_source = torch.tensor([-8, -3, -2])
|
||||
side_target = torch.tensor([1.0, 0.0, 5.0])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "009039_mpii_test":
|
||||
top_source = torch.tensor([-2.0, -3.0, -1.5])
|
||||
top_target = torch.tensor([1.0, 0.0, 3.0])
|
||||
|
||||
side_source = torch.tensor([6, -2, -1.5])
|
||||
side_target = torch.tensor([2.0, 0.0, 6.0])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "009607_mpii_test":
|
||||
top_source = torch.tensor([0, -3, -1])
|
||||
top_target = torch.tensor([0, 0, 5])
|
||||
|
||||
side_source = torch.tensor([-6, -0.5, 2])
|
||||
side_target = torch.tensor([-1.0, 0.5, 7])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "012968_mpii_test":
|
||||
top_source = torch.tensor([-0.5, -3, -4])
|
||||
top_target = torch.tensor([0.0, 0.5, 2.0])
|
||||
|
||||
side_source = torch.tensor([7, -4, 0])
|
||||
side_target = torch.tensor([0, 0.5, 2.0])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "014531_mpii_test":
|
||||
top_source = torch.tensor([0, -2, -2])
|
||||
top_target = torch.tensor([0, 0, 2])
|
||||
|
||||
side_source = torch.tensor([-2, 0, -2])
|
||||
side_target = torch.tensor([0, 0, 2])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "015933_mpii_test":
|
||||
top_source = torch.tensor([0.0, 0.5, -5.0])
|
||||
top_target = torch.tensor([0.0, 2.0, -2.0])
|
||||
|
||||
side_source = torch.tensor([5.0, 1.5, -3.0])
|
||||
side_target = torch.tensor([0.0, 2.5, 1.5])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "016195_mpii_test":
|
||||
top_source = torch.tensor([0, -1, -2])
|
||||
top_target = torch.tensor([0, 1, 1])
|
||||
|
||||
side_source = torch.tensor([-3, 0, 1])
|
||||
side_target = torch.tensor([0, 1, 2])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "018061_mpii_test":
|
||||
top_source = torch.tensor([4, -1, 2])
|
||||
top_target = torch.tensor([1, 1, 3])
|
||||
|
||||
side_source = torch.tensor([-3, 0, 4])
|
||||
side_target = torch.tensor([1, 1, 3])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "018713_mpii_test":
|
||||
top_source = torch.tensor([0, -1, -1])
|
||||
top_target = torch.tensor([0, 0.5, 2])
|
||||
|
||||
side_source = torch.tensor([3, 0, 1])
|
||||
side_target = torch.tensor([0, 0, 3])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "022691_mpii_test":
|
||||
top_source = torch.tensor([0, -4, -4])
|
||||
top_target = torch.tensor([-2.0, 0.0, 1.0])
|
||||
|
||||
side_source = torch.tensor([-5, -2, -5])
|
||||
side_target = torch.tensor([-3.5, 0.0, 1.0])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "023390_mpii_test":
|
||||
top_source = torch.tensor([0.0, -1.5, -1.0])
|
||||
top_target = torch.tensor([0.0, 0.0, 2.0])
|
||||
|
||||
side_source = torch.tensor([4.0, -0.5, 0])
|
||||
side_target = torch.tensor([1.0, 0.5, 4])
|
||||
# top_source = torch.tensor([0.0, -3.0, -2.0])
|
||||
# top_target = torch.tensor([0.0, 0.5, 2.0])
|
||||
|
||||
# side_source = torch.tensor([5.0, -0.5, -3])
|
||||
# side_target = torch.tensor([1.0, 0.5, 4])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name=="024165_mpii_test":
|
||||
top_source = torch.tensor([0, -3, -3])
|
||||
top_target = torch.tensor([0, 0, 3])
|
||||
|
||||
side_source = torch.tensor([-4, -1, -2])
|
||||
side_target = torch.tensor([0, 0, 3])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "024154_mpii_test":
|
||||
top_source = torch.tensor([-1, -0.5, 1])
|
||||
top_target = torch.tensor([-1, 1, 8])
|
||||
|
||||
side_source = torch.tensor([-4, 1, 3])
|
||||
side_target = torch.tensor([-1, 1, 8])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
# if seq_name == "024159_mpii_test":
|
||||
# top_source = torch.tensor([-4.0, -2.0, -6.0])
|
||||
# top_target = torch.tensor([-4.0, 0.0, 0.0])
|
||||
|
||||
# side_source = torch.tensor([-12, -1, -2])
|
||||
# side_target = torch.tensor([-10.0, 0.0, 0.0])
|
||||
# return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "023962_mpii_test":
|
||||
top_source = torch.tensor([0, -0.5, -2])
|
||||
top_target = torch.tensor([0, 0.5, 2])
|
||||
|
||||
side_source = torch.tensor([3, 0.5, 0])
|
||||
side_target = torch.tensor([0, 0, 2])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "017437_mpii_test":
|
||||
top_source = torch.tensor([0, -8, -6])
|
||||
top_target = torch.tensor([-2.0, -1.0, 2.0])
|
||||
|
||||
side_source = torch.tensor([-5, -1, -7])
|
||||
side_target = torch.tensor([-3.0, -1.0, 1.0])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "recording_20211002_S03_S18_04-all-100-200":
|
||||
top_source = torch.tensor([1.0, -1, -4.5])
|
||||
top_target = torch.tensor([1.0, 0.0, 1])
|
||||
|
||||
side_source = torch.tensor([1.0, 0.5, -4.5])
|
||||
side_target = torch.tensor([1.0, 0.5, 1])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "recording_20211002_S03_S18_04-all-100-500":
|
||||
top_source = torch.tensor([0.0, -0.5, -3])
|
||||
top_target = torch.tensor([0.0, 1.0, 1])
|
||||
|
||||
side_source = torch.tensor([-2.0, 0.5, -3])
|
||||
side_target = torch.tensor([0.0, 1.0, 2])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "recording_20211002_S03_S18_01-all-700-862":
|
||||
top_source = torch.tensor([0.5, -2.0, -2])
|
||||
top_target = torch.tensor([0.5, -0.5, 2])
|
||||
|
||||
side_source = torch.tensor([2.0, -0.5, -2])
|
||||
side_target = torch.tensor([0.5, -0.5, 1])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "recording_20220315_S21_S30_02-all-0-100":
|
||||
top_source = torch.tensor([0.0, -1, -2])
|
||||
top_target = torch.tensor([0.0, 0.5, 1])
|
||||
|
||||
side_source = torch.tensor([-2.0, 0.0, -1])
|
||||
side_target = torch.tensor([0.0, 0.5, 1])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if bounds is not None:
|
||||
# (3), (3), (3)
|
||||
bb_min, bb_max, center = bounds
|
||||
print("SCENE BOUNDS", bb_min, bb_max, center)
|
||||
length = torch.abs(bb_max - bb_min).max()
|
||||
print(length)
|
||||
top_source = center + torch.tensor([0.0, -2.0, -0.9 * length])
|
||||
|
||||
side_source = center + torch.tensor([0.5 * length, -0.5, -0.7 * length])
|
||||
return (top_source, center), (side_source, center)
|
||||
|
||||
top_source = torch.tensor([0.0, -2.0, -3.0])
|
||||
top_target = torch.tensor([0.0, 0.0, 1.0])
|
||||
|
||||
side_source = torch.tensor([3.0, -1.0, -1.0])
|
||||
side_target = torch.tensor([0.0, 0.0, 1.0])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
+379
@@ -0,0 +1,379 @@
|
||||
import os
|
||||
import cv2
|
||||
import imageio
|
||||
import numpy as np
|
||||
import torch
|
||||
import subprocess
|
||||
|
||||
from body_model import run_mano
|
||||
from geometry import camera as cam_util
|
||||
from geometry.mesh import make_batch_mesh, save_mesh_scenes
|
||||
from geometry.plane import parse_floor_plane, get_plane_transform
|
||||
|
||||
from util.tensor import detach_all, to_torch, move_to
|
||||
|
||||
from .fig_specs import get_seq_figure_skip, get_seq_static_lookat_points
|
||||
from .tools import mesh_to_geometry, smooth_results
|
||||
|
||||
|
||||
def prep_result_vis(res, vis_mask, track_ids, body_model, temporal_smooth, smooth_trans=True):
|
||||
"""
|
||||
:param res (dict) with (B, T, *) tensor elements, B tracks and T frames
|
||||
:param vis_mask (B, T) with visibility of each track in each frame
|
||||
:param track_ids (B,) index of each track
|
||||
:param temporal_smooth (bool) whether to apply temporal smoothing
|
||||
:param smooth_trans (bool) whether to smooth translation (only used if temporal_smooth=True)
|
||||
"""
|
||||
print("RESULT FIELDS", res.keys())
|
||||
res = detach_all(res)
|
||||
# print(res.keys(), res["is_right"])
|
||||
with torch.no_grad():
|
||||
if temporal_smooth:
|
||||
if smooth_trans:
|
||||
print('running temporal smooth (including trans)')
|
||||
# Smooth everything including trans
|
||||
res["root_orient"], res["pose_body"], res['betas'], res["trans"] = smooth_results(res["root_orient"], res["pose_body"], res['betas'], res["is_right"], res["trans"])
|
||||
else:
|
||||
print('running temporal smooth (excluding trans)')
|
||||
# Smooth only pose and shape, not trans
|
||||
res["root_orient"], res["pose_body"], res['betas'], _ = smooth_results(res["root_orient"], res["pose_body"], res['betas'], res["is_right"])
|
||||
|
||||
world_smpl = run_mano(
|
||||
body_model,
|
||||
res["trans"],
|
||||
res["root_orient"],
|
||||
res["pose_body"],
|
||||
res["is_right"],
|
||||
res.get("betas", None),
|
||||
)
|
||||
|
||||
T_w2c = None
|
||||
floor_plane = None
|
||||
if "cam_R" in res and "cam_t" in res:
|
||||
T_w2c = cam_util.make_4x4_pose(res["cam_R"][0], res["cam_t"][0])
|
||||
if "floor_plane" in res:
|
||||
floor_plane = res["floor_plane"][0]
|
||||
return build_scene_dict(
|
||||
world_smpl,
|
||||
vis_mask,
|
||||
track_ids,
|
||||
T_w2c=T_w2c,
|
||||
floor_plane=floor_plane,
|
||||
)
|
||||
|
||||
|
||||
def build_scene_dict(
|
||||
world_smpl, vis_mask, track_ids, T_w2c=None, floor_plane=None, **kwargs
|
||||
):
|
||||
scene_dict = {}
|
||||
|
||||
# first get the geometry of the people
|
||||
# lists of length T with (B, V, 3), (F, 3), (B, 3)
|
||||
scene_dict["geometry"] = mesh_to_geometry(
|
||||
world_smpl["vertices"], world_smpl["joints"], world_smpl["l_faces"], world_smpl["r_faces"], world_smpl["is_right"], vis_mask, track_ids
|
||||
)
|
||||
|
||||
if T_w2c is None:
|
||||
T_w2c = torch.eye(4)[None]
|
||||
|
||||
T_c2w = torch.linalg.inv(T_w2c)
|
||||
# rotate the camera slightly down and translate back and up
|
||||
T = cam_util.make_4x4_pose(
|
||||
cam_util.rotx(-np.pi / 10), torch.tensor([0, -1, -2])
|
||||
).to(T_c2w.device)
|
||||
|
||||
scene_dict["cameras"] = {
|
||||
"src_cam": T_c2w,
|
||||
# "front": torch.einsum("ij,...jk->...ik", T, T_c2w),
|
||||
}
|
||||
|
||||
# Create ground plane based on hand mesh height
|
||||
# DISABLED: Ground plane visualization disabled to better see coordinate system
|
||||
# verts = scene_dict["geometry"][0] # Get vertices from geometry
|
||||
# if len(verts) > 0:
|
||||
# # Find minimum height across all frames and vertices
|
||||
# min_height = float('inf')
|
||||
# for frame_verts in verts:
|
||||
# if len(frame_verts) > 0:
|
||||
# frame_min = frame_verts[..., 1].min().item() # y-coordinate is height
|
||||
# min_height = min(min_height, frame_min)
|
||||
#
|
||||
# # Set ground plane further below minimum height
|
||||
# ground_offset = -0.5 # 20cm below minimum height
|
||||
# ground_height = min_height - ground_offset
|
||||
#
|
||||
# # Create ground plane transform
|
||||
# R = torch.eye(3) # Identity rotation (flat ground)
|
||||
# t = torch.tensor([0.0, ground_height, 0.0]) # Translate to ground height
|
||||
# scene_dict["ground"] = cam_util.make_4x4_pose(R, t)
|
||||
#
|
||||
# # Save ground mesh for Blender debugging
|
||||
# import trimesh
|
||||
# from vis.viewer import make_checkerboard
|
||||
# ground_mesh = make_checkerboard(color0=[0.9, 0.95, 1.0], color1=[0.7, 0.8, 0.85], up="y", alpha=1.0)
|
||||
# ground_mesh.apply_translation([0.0, ground_height, 0.0])
|
||||
# ground_mesh.export("ground_debug.obj")
|
||||
|
||||
# if floor_plane is not None:
|
||||
# # compute the ground transform
|
||||
# # use the first appearance of a track as the reference point
|
||||
# tid, sid = torch.where(vis_mask > 0)
|
||||
# idx = tid[torch.argmin(sid)]
|
||||
# root = world_smpl["joints"][idx, 0, 0].detach().cpu()
|
||||
# floor = parse_floor_plane(floor_plane.detach().cpu())
|
||||
# R, t = get_plane_transform(torch.tensor([0.0, 1.0, 0.0]), floor, root)
|
||||
# scene_dict["ground"] = cam_util.make_4x4_pose(R, t)
|
||||
|
||||
return scene_dict
|
||||
|
||||
|
||||
# def render_scene_dict(renderer, scene_dict, out_name, fps=30, **kwargs):
|
||||
# # lists of T (B, V, 3), (B, 3), (F, 3)
|
||||
# verts, colors, faces, bounds = scene_dict["geometry"]
|
||||
# print("NUM VERTS", len(verts))
|
||||
|
||||
# # add a top view
|
||||
# scene_dict["cameras"]["above"] = cam_util.make_4x4_pose(
|
||||
# torch.eye(3), torch.tensor([0, 0, -10])
|
||||
# )[None]
|
||||
|
||||
# for cam_name, cam_poses in scene_dict["cameras"].items():
|
||||
# print("rendering scene for", cam_name)
|
||||
# # cam_poses are (T, 4, 4)
|
||||
# render_bg = cam_name == "src_cam"
|
||||
# ground_pose = scene_dict.get("ground", None)
|
||||
# print(cam_name, ground_pose)
|
||||
# frames = renderer.render_video(
|
||||
# cam_poses[None], verts, faces, colors, render_bg, ground_pose=ground_pose
|
||||
# )
|
||||
# os.makedirs(f"{out_name}_{cam_name}/", exist_ok=True)
|
||||
# for idx, i in enumerate(frames):
|
||||
# imageio.imwrite(f"{out_name}_{cam_name}/" + f'{str(idx).zfill(6)}.jpg', i)
|
||||
# imageio.mimwrite(f"{out_name}_{cam_name}.mp4", frames, fps=fps)
|
||||
|
||||
|
||||
def animate_scene(
|
||||
vis,
|
||||
scene,
|
||||
out_name,
|
||||
seq_name=None,
|
||||
accumulate=False,
|
||||
render_views=["src_cam", "front", "above", "side"],
|
||||
render_bg=True,
|
||||
render_cam=True,
|
||||
render_ground=True,
|
||||
debug=False,
|
||||
render_keypoints=False,
|
||||
**kwargs,
|
||||
):
|
||||
if len(render_views) < 1:
|
||||
return
|
||||
|
||||
scene = build_pyrender_scene(
|
||||
vis,
|
||||
scene,
|
||||
seq_name,
|
||||
render_views=render_views,
|
||||
render_cam=render_cam,
|
||||
accumulate=accumulate,
|
||||
debug=debug,
|
||||
)
|
||||
|
||||
print("RENDERING VIEWS", scene["cameras"].keys())
|
||||
render_ground = render_ground and "ground" in scene
|
||||
save_paths = []
|
||||
for cam_name, cam_poses in scene["cameras"].items():
|
||||
is_src = cam_name == "src_cam"
|
||||
show_bg = is_src and render_bg
|
||||
show_ground = render_ground and not is_src
|
||||
show_cam = render_cam and not is_src
|
||||
show_keypoints = is_src and render_keypoints # Only show keypoints on source view
|
||||
vis_name = f"{out_name}_{cam_name}"
|
||||
print(f"{cam_name} has {len(cam_poses)} poses")
|
||||
skip = 10 if debug else 1
|
||||
vis.set_camera_seq(cam_poses[::skip])
|
||||
save_path = vis.animate(
|
||||
vis_name,
|
||||
render_bg=show_bg,
|
||||
render_ground=show_ground,
|
||||
render_cam=show_cam,
|
||||
render_keypoints=show_keypoints,
|
||||
**kwargs,
|
||||
)
|
||||
save_paths.append(save_path)
|
||||
|
||||
return save_paths
|
||||
|
||||
|
||||
def build_pyrender_scene(
|
||||
vis,
|
||||
scene,
|
||||
seq_name,
|
||||
render_views=["src_cam", "front", "above", "side"],
|
||||
render_cam=True,
|
||||
accumulate=False,
|
||||
debug=False,
|
||||
):
|
||||
"""
|
||||
:param vis (viewer object)
|
||||
:param scene (scene dict with geometry, cameras, etc)
|
||||
:param accumulate (optional bool, default False) whether to render entire trajectory together
|
||||
:param render_views (list str) camera views to render
|
||||
"""
|
||||
if len(render_views) < 1:
|
||||
return
|
||||
|
||||
assert all(view in ["src_cam", "front", "above", "side"] for view in render_views)
|
||||
|
||||
scene = move_to(detach_all(scene), "cpu")
|
||||
src_cams = scene["cameras"]["src_cam"]
|
||||
verts, _, colors, l_faces, r_faces, is_right, bounds = scene["geometry"]
|
||||
T = len(verts)
|
||||
print(f"{T} mesh frames for {seq_name}, {len(verts)}")
|
||||
|
||||
# set camera views
|
||||
if not "cameras" in scene:
|
||||
scene["cameras"] = {}
|
||||
|
||||
# remove default views from source camera perspective if desired
|
||||
if "src_cam" not in render_views:
|
||||
scene["cameras"].pop("src_cam", None)
|
||||
if "front" not in render_views:
|
||||
scene["cameras"].pop("front", None)
|
||||
|
||||
# add static viewpoints if desired
|
||||
top_pose, side_pose, _skip = get_static_views(seq_name, bounds)
|
||||
if "above" in render_views:
|
||||
scene["cameras"]["above"] = top_pose[None]
|
||||
if "side" in render_views:
|
||||
scene["cameras"]["side"] = side_pose[None]
|
||||
if "front" in render_views:
|
||||
# Use a static front camera: place it in front of the hand mesh, looking at the center
|
||||
# Place the camera at a fixed distance along the -z axis from the center
|
||||
if bounds is not None:
|
||||
bb_min, bb_max, center = bounds
|
||||
length = torch.abs(bb_max - bb_min).max()
|
||||
front_offset = 1.0 # 1 closer, 2 farther
|
||||
front_source = center + torch.tensor([0.0, -0.5, -front_offset])
|
||||
front_target = center
|
||||
up = torch.tensor([0.0, 1, 0.0])
|
||||
front_pose = cam_util.lookat_matrix(front_source, front_target, up)
|
||||
scene["cameras"]["front"] = front_pose[None]
|
||||
else:
|
||||
# fallback to previous behavior if bounds not available
|
||||
pass
|
||||
|
||||
# accumulate meshes if possible (can only accumulate for static camera)
|
||||
moving_cam = "src_cam" in render_views or "front" in render_views
|
||||
accumulate = accumulate and not moving_cam
|
||||
skip = _skip if accumulate else 1
|
||||
|
||||
vis.clear_meshes()
|
||||
|
||||
if "ground" in scene:
|
||||
vis.set_ground(scene["ground"])
|
||||
|
||||
if debug:
|
||||
skip = 10
|
||||
times = list(range(0, T, skip))
|
||||
# print(times, verts[0].shape)
|
||||
# raise ValueError
|
||||
for t in times:
|
||||
if len(is_right[t]) > 1:
|
||||
assert (is_right[t].cpu().numpy().tolist() == [0,1])
|
||||
l_meshes = make_batch_mesh(verts[t][0][None], l_faces[t], colors[t][0][None])
|
||||
r_meshes = make_batch_mesh(verts[t][1][None], r_faces[t], colors[t][1][None])
|
||||
assert len(l_meshes) == 1
|
||||
assert len(r_meshes) == 1
|
||||
meshes = [l_meshes[0], r_meshes[0]]
|
||||
else:
|
||||
assert len(is_right[t]) == 1
|
||||
if is_right[t] == 0:
|
||||
meshes = make_batch_mesh(verts[t][0][None], l_faces[t], colors[t][0][None])
|
||||
elif is_right[t] == 1:
|
||||
meshes = make_batch_mesh(verts[t][0][None], r_faces[t], colors[t][0][None])
|
||||
|
||||
if accumulate:
|
||||
vis.add_static_meshes(meshes)
|
||||
else:
|
||||
vis.add_mesh_frame(meshes, debug=debug)
|
||||
|
||||
# add camera markers
|
||||
if render_cam:
|
||||
if accumulate:
|
||||
vis.add_camera_markers_static(src_cams[::skip])
|
||||
else:
|
||||
vis.add_camera_markers(src_cams[::skip])
|
||||
|
||||
return scene
|
||||
|
||||
|
||||
def get_static_views(seq_name=None, bounds=None):
|
||||
print("STATIC VIEWS FOR SEQ NAME", seq_name)
|
||||
up = torch.tensor([0.0, 1.0, 0.0])
|
||||
|
||||
skip = get_seq_figure_skip(seq_name)
|
||||
top_vp, side_vp = get_seq_static_lookat_points(seq_name, bounds)
|
||||
top_source, top_target = top_vp
|
||||
side_source, side_target = side_vp
|
||||
top_pose = cam_util.lookat_matrix(top_source, top_target, up)
|
||||
side_pose = cam_util.lookat_matrix(side_source, side_target, up)
|
||||
return top_pose, side_pose, skip
|
||||
|
||||
|
||||
def make_video_grid_2x2(out_path, vid_paths, fps, overwrite=False):
|
||||
if os.path.isfile(out_path) and not overwrite:
|
||||
print(f"{out_path} already exists, skipping.")
|
||||
return
|
||||
|
||||
if any(not os.path.isfile(v) for v in vid_paths):
|
||||
print("not all inputs exist!", vid_paths)
|
||||
return
|
||||
|
||||
# Open all videos for reading
|
||||
readers = [imageio.get_reader(v, 'ffmpeg') for v in vid_paths]
|
||||
try:
|
||||
# Check min frame count to avoid index errors; also get shape/fps
|
||||
lengths = [reader.count_frames() for reader in readers]
|
||||
min_len = min(lengths)
|
||||
if min_len == 0:
|
||||
print("Some videos have 0 frames!")
|
||||
for reader in readers:
|
||||
reader.close()
|
||||
return
|
||||
|
||||
# We assume the first video dictates the output fps and (half)size
|
||||
meta = readers[0].get_meta_data()
|
||||
# fps = meta.get('fps', 30)
|
||||
orig_size = meta['size'] # (width, height)
|
||||
width = orig_size[0] // 2
|
||||
height = orig_size[1] // 2
|
||||
out_size = (width * 2, height * 2)
|
||||
|
||||
writer = imageio.get_writer(out_path, fps=fps, codec='libx264', ffmpeg_log_level="quiet", macro_block_size=None)
|
||||
|
||||
for t in range(min_len):
|
||||
frames = []
|
||||
for r in readers:
|
||||
try:
|
||||
frame = r.get_data(t)
|
||||
except Exception:
|
||||
frame = None
|
||||
frames.append(frame)
|
||||
if any(f is None for f in frames):
|
||||
break
|
||||
# Resize all to (width, height)
|
||||
frames = [cv2.resize(f, (width, height)) for f in frames]
|
||||
row1 = np.hstack(frames[:2])
|
||||
row2 = np.hstack(frames[2:])
|
||||
grid = np.vstack([row1, row2])
|
||||
# Ensure uint8
|
||||
if grid.dtype != np.uint8:
|
||||
grid = np.clip(grid, 0, 255).astype(np.uint8)
|
||||
writer.append_data(grid)
|
||||
|
||||
writer.close()
|
||||
print(f"Video grid created: {out_path}")
|
||||
finally:
|
||||
for reader in readers:
|
||||
reader.close()
|
||||
+254
@@ -0,0 +1,254 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import numpy as np
|
||||
|
||||
from pytorch3d.renderer import (
|
||||
RasterizationSettings,
|
||||
MeshRenderer,
|
||||
MeshRasterizer,
|
||||
HardPhongShader,
|
||||
PointLights,
|
||||
PerspectiveCameras,
|
||||
)
|
||||
from pytorch3d.structures import Meshes
|
||||
from pytorch3d.structures.meshes import join_meshes_as_scene
|
||||
from pytorch3d.renderer import TexturesAtlas, TexturesVertex
|
||||
from pytorch3d.transforms import euler_angles_to_matrix
|
||||
|
||||
from .tools import get_colors, read_image, transform_torch3d, checkerboard_geometry
|
||||
|
||||
|
||||
def init_renderer(img_size, intrins, device, vis_scale=1.0, bg_paths=None):
|
||||
img_size = int(vis_scale * img_size[0]), int(vis_scale * img_size[1])
|
||||
renderer = RenderBase(device, img_size, intrins=vis_scale * intrins)
|
||||
if bg_paths is not None:
|
||||
bg_imgs = [
|
||||
torch.from_numpy(read_image(p, scale=vis_scale)).float() / 255
|
||||
for p in bg_paths
|
||||
]
|
||||
print(bg_imgs[0].shape)
|
||||
renderer.set_bg_seq(bg_imgs)
|
||||
print("RENDERER", renderer)
|
||||
return renderer
|
||||
|
||||
|
||||
class RenderBase(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
device,
|
||||
img_size,
|
||||
intrins=None,
|
||||
blur_radius: float = 0.0, # rasterizer settings
|
||||
faces_per_pixel: int = 3,
|
||||
):
|
||||
"""
|
||||
Uses Pytorch3D to render results.
|
||||
We use the PyTorch3D coordinate system, which assumes:
|
||||
+X:left, +Y: up and +Z: from us to scene (right-handed)
|
||||
https://pytorch3d.org/docs/cameras
|
||||
"""
|
||||
super().__init__()
|
||||
self.device = device
|
||||
W, H = img_size
|
||||
self.img_size = (H, W)
|
||||
if intrins is not None:
|
||||
intrins = intrins.to(device)
|
||||
self.cam_f, self.cam_center = intrins[:2], intrins[2:]
|
||||
cx, cy = self.cam_center.detach().cpu()
|
||||
else:
|
||||
self.cam_f = torch.tensor([0.5 * (H + W), 0.5 * (H + W)], device=device)
|
||||
self.cam_center = torch.tensor([0.5 * W, 0.5 * H], device=device)
|
||||
print("cam_f", self.cam_f, "cam_center", self.cam_center)
|
||||
print("img_size", self.img_size)
|
||||
self.blur_radius = blur_radius
|
||||
self.faces_per_pixel = faces_per_pixel
|
||||
|
||||
self.light_loc = [[0, 0, -0.5]]
|
||||
self.cameras = self.make_cameras()
|
||||
self.bg_imgs = None
|
||||
self.set_ground()
|
||||
self.build_renderer()
|
||||
|
||||
def build_renderer(self):
|
||||
self.lights = PointLights(device=self.device, location=self.light_loc)
|
||||
|
||||
self.rasterizer = RasterizationSettings(
|
||||
image_size=self.img_size,
|
||||
blur_radius=self.blur_radius,
|
||||
faces_per_pixel=self.faces_per_pixel,
|
||||
max_faces_per_bin=100000,
|
||||
)
|
||||
self.renderer = MeshRenderer(
|
||||
rasterizer=MeshRasterizer(
|
||||
cameras=self.cameras, raster_settings=self.rasterizer
|
||||
),
|
||||
shader=HardPhongShader(
|
||||
cameras=self.cameras, lights=self.lights, device=self.device
|
||||
),
|
||||
)
|
||||
|
||||
def make_cameras(self, cam_poses=None):
|
||||
device = self.device
|
||||
if cam_poses is None:
|
||||
cam_poses = torch.eye(4, device=device)[None]
|
||||
cam_poses = cam_poses.to(device)
|
||||
cam_R, cam_t = transform_torch3d(cam_poses)
|
||||
Nc = len(cam_poses)
|
||||
focal = self.cam_f[None].repeat(Nc, 1)
|
||||
ppt = self.cam_center[None].repeat(Nc, 1)
|
||||
img_size = [self.img_size for _ in range(Nc)]
|
||||
return PerspectiveCameras(
|
||||
device=self.device,
|
||||
in_ndc=False,
|
||||
focal_length=focal,
|
||||
principal_point=ppt,
|
||||
image_size=img_size,
|
||||
R=cam_R,
|
||||
T=cam_t,
|
||||
)
|
||||
|
||||
def render_frame(self, cam_poses, verts, faces, colors):
|
||||
"""
|
||||
:param cam_poses (C, 4, 4)
|
||||
:param verts (B, V, 3)
|
||||
:param faces (F, 3)
|
||||
:param colors (B, 4)
|
||||
"""
|
||||
with torch.no_grad():
|
||||
cameras = self.make_cameras(cam_poses)
|
||||
mesh = build_meshes(verts, faces, colors)
|
||||
image = self.renderer(mesh, cameras=cameras)
|
||||
return image
|
||||
|
||||
def render_with_ground(self, cam_poses, verts, faces, colors, ground_pose):
|
||||
"""
|
||||
:param cam_poses (C, 4, 4)
|
||||
:param verts (B, V, 3)
|
||||
:param faces (F, 3)
|
||||
:param colors (B, 3)
|
||||
:param ground_pose (4, 4)
|
||||
"""
|
||||
with torch.no_grad():
|
||||
self.set_ground_pose(ground_pose)
|
||||
cameras = self.make_cameras(cam_poses)
|
||||
# (B, V, 3), (B, F, 3), (B, V, 3)
|
||||
verts, faces, colors = prep_shared_geometry(verts, faces, colors)
|
||||
# (V, 3), (F, 3), (V, 3)
|
||||
gv, gf, gc = self.ground_geometry
|
||||
verts = list(torch.unbind(verts, dim=0)) + [gv]
|
||||
faces = list(torch.unbind(faces, dim=0)) + [gf]
|
||||
colors = list(torch.unbind(colors, dim=0)) + [gc[..., :3]]
|
||||
mesh = create_meshes(verts, faces, colors)
|
||||
|
||||
image = self.renderer(mesh, cameras=cameras)
|
||||
return image
|
||||
|
||||
def set_bg_seq(self, bg_imgs):
|
||||
self.bg_imgs = bg_imgs
|
||||
|
||||
def set_ground(self, pose=None):
|
||||
device = self.device
|
||||
v, f, vc, fc = map(torch.from_numpy, checkerboard_geometry(up="y"))
|
||||
v, f, vc = v.to(device), f.to(device), vc.to(device)
|
||||
self.ground_geometry = [v, f, vc]
|
||||
if pose is not None:
|
||||
self.set_ground_pose(pose)
|
||||
|
||||
def set_ground_pose(self, pose):
|
||||
"""
|
||||
:param pose (4, 4)
|
||||
"""
|
||||
if self.ground_geometry is None:
|
||||
return
|
||||
v = self.ground_geometry[0]
|
||||
pose = pose.to(self.device)
|
||||
# R, t = pose[:3, :3], pose[:3, 3]
|
||||
R, t = transform_torch3d(pose)
|
||||
self.ground_geometry[0] = torch.einsum("ij,vj->vi", R, v) + t[None]
|
||||
|
||||
def composite_bg(self, frames, fac=1.0):
|
||||
T = len(frames)
|
||||
if self.bg_imgs is None:
|
||||
return frames
|
||||
|
||||
out_frames = []
|
||||
for bg, frame in zip(self.bg_imgs, frames):
|
||||
alpha = fac * frame[..., 3:]
|
||||
out = alpha * frame[..., :3] + (1 - alpha) * bg
|
||||
out_frames.append(out)
|
||||
return out_frames
|
||||
|
||||
def render_video(
|
||||
self,
|
||||
cam_poses,
|
||||
verts,
|
||||
faces,
|
||||
colors,
|
||||
render_bg=False,
|
||||
fac=1.0,
|
||||
ground_pose=None,
|
||||
):
|
||||
"""
|
||||
:param cam_poses (C, T, 4, 4)
|
||||
:param verts list of T (B, V, 3)
|
||||
:param faces list of T (F, 3)
|
||||
:param colors list of T (B, 4)
|
||||
:param render_bg (optional default False)
|
||||
:param fac (optional float, default 1)
|
||||
:param ground_pose (optional, default None) (4, 4)
|
||||
returns list of T (H, W, 4)
|
||||
"""
|
||||
T = len(verts)
|
||||
print(f"rendering video with {T} frames")
|
||||
if cam_poses.shape[1] == 1:
|
||||
cam_poses = cam_poses.expand(-1, T, -1, -1)
|
||||
|
||||
if render_bg: # don't render ground on source frame
|
||||
ground_pose = None
|
||||
|
||||
frames = []
|
||||
for t in range(T):
|
||||
if ground_pose is None:
|
||||
# (H, W, 4)
|
||||
frame = self.render_frame(
|
||||
cam_poses[:, t], verts[t], faces[t], colors[t]
|
||||
)
|
||||
else:
|
||||
# (H, W, 4)
|
||||
frame = self.render_with_ground(
|
||||
cam_poses[:, t], verts[t], faces[t], colors[t], ground_pose
|
||||
)
|
||||
frames.append(frame.cpu()[0])
|
||||
|
||||
if render_bg and self.bg_imgs is not None:
|
||||
frames = self.composite_bg(frames, fac=fac)
|
||||
|
||||
return [(255 * f).byte() for f in frames]
|
||||
|
||||
|
||||
def build_meshes(verts, faces, colors):
|
||||
return create_meshes(*prep_shared_geometry(verts, faces, colors))
|
||||
|
||||
|
||||
def prep_shared_geometry(verts, faces, colors):
|
||||
"""
|
||||
:param verts (B, V, 3)
|
||||
:param faces (F, 3)
|
||||
:param colors (B, 4)
|
||||
"""
|
||||
B, V, _ = verts.shape
|
||||
F, _ = faces.shape
|
||||
colors = colors.unsqueeze(1).expand(B, V, -1)[..., :3]
|
||||
faces = faces.unsqueeze(0).expand(B, F, -1)
|
||||
return verts, faces, colors
|
||||
|
||||
|
||||
def create_meshes(verts, faces, colors):
|
||||
"""
|
||||
:param verts (B, V, 3)
|
||||
:param faces (B, F, 3)
|
||||
:param colors (B, V, 3)
|
||||
"""
|
||||
textures = TexturesVertex(verts_features=colors)
|
||||
meshes = Meshes(verts=verts, faces=faces, textures=textures)
|
||||
return join_meshes_as_scene(meshes)
|
||||
+1109
File diff suppressed because it is too large
Load Diff
+672
@@ -0,0 +1,672 @@
|
||||
import os
|
||||
import imageio
|
||||
import numpy as np
|
||||
|
||||
import time
|
||||
import torch
|
||||
import trimesh
|
||||
|
||||
os.environ["PYOPENGL_PLATFORM"] = "egl"
|
||||
|
||||
import pyrender
|
||||
from pyrender.constants import RenderFlags
|
||||
from pyrender.light import DirectionalLight
|
||||
from pyrender.node import Node
|
||||
|
||||
from geometry.camera import make_rotation, make_translation
|
||||
from .tools import (
|
||||
read_image,
|
||||
checkerboard_geometry,
|
||||
camera_marker_geometry,
|
||||
transform_pyrender,
|
||||
)
|
||||
|
||||
|
||||
def init_viewer(
|
||||
img_size, intrins, vis_scale=1.0, bg_paths=None, fps=30,
|
||||
):
|
||||
img_size = int(vis_scale * img_size[0]), int(vis_scale * img_size[1])
|
||||
intrins = vis_scale * intrins
|
||||
|
||||
platform = os.environ.get("PYOPENGL_PLATFORM", "pyglet")
|
||||
if platform == "pyglet":
|
||||
vis = AnimationViewer(img_size, intrins=intrins, fps=fps)
|
||||
print("VIS", vis)
|
||||
return vis
|
||||
|
||||
if platform != "egl":
|
||||
raise NotImplementedError
|
||||
|
||||
vis = OffscreenAnimation(img_size, intrins=intrins, fps=fps)
|
||||
if bg_paths is not None:
|
||||
bg_imgs = [
|
||||
read_image(p, scale=vis_scale).astype(np.float32) / 255 for p in bg_paths
|
||||
]
|
||||
vis.set_bg_seq(bg_imgs)
|
||||
print("VIS", vis)
|
||||
return vis
|
||||
|
||||
|
||||
class AnimationBase(object):
|
||||
"""
|
||||
Render an animation
|
||||
IMPORTANT: pyrender uses coordinates with Y UP (i.e. xyz = right up back),
|
||||
we must keep this in mind because the human prior frame has Z UP (i.e.
|
||||
xyz = right forward up)
|
||||
"""
|
||||
|
||||
def __init__(self, img_size, intrins=None):
|
||||
self.scene = pyrender.Scene(
|
||||
ambient_light=[0.3, 0.3, 0.3], bg_color=[1.0, 1.0, 1.0, 0.0]
|
||||
)
|
||||
self.scene.bg_color = np.array([1.0, 1.0, 1.0, 0.0])
|
||||
|
||||
self.viewport_size = img_size
|
||||
self.camera = make_pyrender_camera(img_size, intrins)
|
||||
|
||||
# set the ground with Z up
|
||||
self.ground_pose = np.eye(4)
|
||||
ground = pyrender.Mesh.from_trimesh(
|
||||
make_checkerboard(color0=[0.9, 0.95, 1.0], color1=[0.7, 0.8, 0.85], up="y", alpha=1.0), smooth=False
|
||||
)
|
||||
self.ground_node = self.scene.add(ground, name="ground", pose=self.ground_pose)
|
||||
self.cam_node = self.scene.add(self.camera, name="camera")
|
||||
|
||||
self.cam_marker_node = None
|
||||
self.cam_marker_poses = []
|
||||
|
||||
self.light_nodes = []
|
||||
self.light_poses = []
|
||||
|
||||
self.bg_seq = []
|
||||
|
||||
# pyrender nodes, reused between frames
|
||||
self.anim_nodes = []
|
||||
# list of meshes per frame of animation
|
||||
self.anim_meshes = []
|
||||
# list of camera poses per frame of animation
|
||||
self.anim_cameras = None
|
||||
# current timestep of animation
|
||||
self.anim_idx = 0
|
||||
|
||||
def acquire_lock(self):
|
||||
pass
|
||||
|
||||
def release_lock(self):
|
||||
pass
|
||||
|
||||
def close(self):
|
||||
pass
|
||||
|
||||
def delete(self):
|
||||
pass
|
||||
|
||||
def clear_meshes(self):
|
||||
self.anim_meshes = []
|
||||
self.anim_idx = 0
|
||||
|
||||
@property
|
||||
def anim_len(self):
|
||||
return len(self.anim_meshes)
|
||||
|
||||
def add_lighting(self, color=np.ones(3), intensity=1.0):
|
||||
self.light_poses = get_light_poses()
|
||||
self.light_poses.append(np.eye(4))
|
||||
cam_pose = self.scene.get_pose(self.cam_node)
|
||||
for i, pose in enumerate(self.light_poses):
|
||||
matrix = cam_pose @ pose
|
||||
node = Node(
|
||||
name=f"light-{i:02d}",
|
||||
light=DirectionalLight(color=color, intensity=intensity),
|
||||
matrix=matrix,
|
||||
)
|
||||
if self.scene.has_node(node):
|
||||
continue
|
||||
self.scene.add_node(node)
|
||||
self.light_nodes.append(node)
|
||||
|
||||
def add_mesh_frame(self, meshes, debug=False):
|
||||
"""
|
||||
:param meshes (list) trimesh meshes for this frame timestep
|
||||
"""
|
||||
meshes = [pyrender.Mesh.from_trimesh(m) for m in meshes]
|
||||
num_nodes = len(meshes)
|
||||
self.anim_meshes.append(meshes)
|
||||
|
||||
if num_nodes > len(self.anim_nodes):
|
||||
# need to add a node in the pyrender scene
|
||||
self.acquire_lock()
|
||||
for i in range(len(self.anim_nodes), num_nodes):
|
||||
self.anim_nodes.append(self.scene.add(meshes[i], name=f"mesh_{i:03d}"))
|
||||
self.release_lock()
|
||||
|
||||
def add_static_meshes(self, meshes, smooth=True):
|
||||
"""
|
||||
add all meshes to a single frame (timestep)
|
||||
"""
|
||||
meshes = [pyrender.Mesh.from_trimesh(m, smooth=smooth) for m in meshes]
|
||||
if len(self.anim_meshes) < 1:
|
||||
self.anim_meshes.append([])
|
||||
self.anim_meshes[0].extend(meshes)
|
||||
num_meshes = len(self.anim_meshes[0])
|
||||
num_nodes = len(self.anim_nodes)
|
||||
if num_meshes > num_nodes:
|
||||
# need to add a node in the pyrender scene
|
||||
self.acquire_lock()
|
||||
self.anim_nodes.extend(
|
||||
[
|
||||
self.scene.add(self.anim_meshes[0][i])
|
||||
for i in range(num_nodes, num_meshes)
|
||||
]
|
||||
)
|
||||
self.release_lock()
|
||||
|
||||
def set_camera_seq(self, poses):
|
||||
"""
|
||||
:param poses (*, 4, 4)
|
||||
"""
|
||||
poses = transform_pyrender(poses)
|
||||
poses = np.asarray(poses)
|
||||
assert poses.ndim >= 2 and poses.shape[-2:] == (4, 4)
|
||||
if poses.ndim < 3:
|
||||
poses = poses[None]
|
||||
print(f"Adding camera sequence length {len(poses)}")
|
||||
self.anim_cameras = poses
|
||||
|
||||
def add_camera_markers(self, poses):
|
||||
"""
|
||||
Add a single camera marker node that we move around when we animate
|
||||
"""
|
||||
print("ADDING CAMERA MARKERS")
|
||||
if self.cam_marker_node is None:
|
||||
cam_marker = make_camera_marker(up="y")
|
||||
self.cam_marker_node = self.scene.add(
|
||||
pyrender.Mesh.from_trimesh(cam_marker, smooth=False)
|
||||
)
|
||||
poses = transform_pyrender(poses)
|
||||
poses = np.asarray(poses)
|
||||
self.cam_marker_poses = poses
|
||||
|
||||
def add_camera_markers_static(self, poses):
|
||||
"""
|
||||
Add all camera markers as separate mesh nodes to render in one pass
|
||||
"""
|
||||
print("ADDING STATIC CAMERA MARKERS")
|
||||
poses = transform_pyrender(poses)
|
||||
poses = np.asarray(poses)
|
||||
cam_meshes = [make_camera_marker(up="y", transform=pose) for pose in poses]
|
||||
self.add_static_meshes(cam_meshes, smooth=False)
|
||||
|
||||
def update_camera(self, cam_pose):
|
||||
self.acquire_lock()
|
||||
self.scene.set_pose(self.cam_node, pose=cam_pose)
|
||||
self.release_lock()
|
||||
|
||||
def set_ground(self, pose):
|
||||
pose = transform_pyrender(pose)
|
||||
print("Setting ground pose to be", pose)
|
||||
self.ground_pose = np.asarray(pose)
|
||||
self.scene.set_pose(self.ground_node, pose=self.ground_pose)
|
||||
|
||||
def set_bg_seq(self, bg_imgs):
|
||||
pass
|
||||
|
||||
def set_mesh_visibility(self, vis):
|
||||
self.acquire_lock()
|
||||
for node in self.anim_nodes:
|
||||
node.mesh.is_visible = vis
|
||||
self.release_lock()
|
||||
|
||||
def check_mesh_visibility(self):
|
||||
for node in self.anim_nodes:
|
||||
if node.mesh.is_visible:
|
||||
return True
|
||||
|
||||
if self.ground_node.mesh.is_visible:
|
||||
return True
|
||||
|
||||
if self.cam_marker_node is not None and self.cam_marker_node.mesh.is_visible:
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
def update_frame(self):
|
||||
t = self.anim_idx % self.anim_len
|
||||
|
||||
# update camera
|
||||
if self.anim_cameras is not None:
|
||||
cam_t = min(t, len(self.anim_cameras) - 1)
|
||||
cam_pose = self.anim_cameras[cam_t]
|
||||
self.update_camera(cam_pose)
|
||||
if self.cam_marker_node is not None and t < len(self.cam_marker_poses):
|
||||
self.acquire_lock()
|
||||
self.scene.set_pose(self.cam_marker_node, self.cam_marker_poses[t])
|
||||
self.release_lock()
|
||||
|
||||
# update meshes
|
||||
meshes = self.anim_meshes[t]
|
||||
self.acquire_lock()
|
||||
for i, node in enumerate(self.anim_nodes):
|
||||
if i < len(meshes):
|
||||
node.mesh = meshes[i]
|
||||
node.mesh.is_visible = True
|
||||
else:
|
||||
node.mesh.is_visible = False
|
||||
self.release_lock()
|
||||
|
||||
|
||||
class OffscreenAnimation(AnimationBase):
|
||||
def __init__(self, img_size, intrins=None, fps=30, ext="mp4"):
|
||||
super().__init__(img_size, intrins=intrins)
|
||||
self.add_lighting(0.9)
|
||||
|
||||
self.fps = fps
|
||||
self.ext = ext
|
||||
self.viewer = pyrender.OffscreenRenderer(*self.viewport_size)
|
||||
|
||||
self.bg_seq = []
|
||||
self.keypoints_seq = None # Store GT 2D keypoints for overlay
|
||||
self.pred_keypoints_seq = None # Store predicted 2D keypoints
|
||||
|
||||
def set_bg_seq(self, bg_imgs):
|
||||
assert isinstance(bg_imgs, list)
|
||||
print(f"setting length {len(bg_imgs)} bg sequence")
|
||||
self.bg_seq = bg_imgs
|
||||
|
||||
def set_keypoints_seq(self, keypoints_seq):
|
||||
"""
|
||||
Set GT 2D keypoints sequence for overlay visualization
|
||||
Args:
|
||||
keypoints_seq: List of (J, 3) arrays with [x, y, confidence] for each frame
|
||||
"""
|
||||
self.keypoints_seq = keypoints_seq
|
||||
print(f"setting length {len(keypoints_seq)} GT keypoints sequence")
|
||||
|
||||
def set_pred_keypoints_seq(self, pred_keypoints_seq):
|
||||
"""
|
||||
Set predicted 2D keypoints sequence
|
||||
Args:
|
||||
pred_keypoints_seq: List of (J, 2) arrays with [x, y] for each frame
|
||||
"""
|
||||
self.pred_keypoints_seq = pred_keypoints_seq
|
||||
print(f"setting length {len(pred_keypoints_seq)} predicted keypoints sequence")
|
||||
|
||||
def delete(self):
|
||||
self.viewer.delete()
|
||||
|
||||
def close(self):
|
||||
self.delete()
|
||||
|
||||
def draw_keypoints_on_image(self, img, keypoints, radius=4, color=(0, 255, 0), thickness=-1, conf_threshold=0.3):
|
||||
"""
|
||||
Draw 2D keypoints on an image
|
||||
Args:
|
||||
img: numpy array (H, W, 3)
|
||||
keypoints: (J, 3) array with [x, y, confidence]
|
||||
radius: circle radius for keypoints
|
||||
color: BGR color tuple
|
||||
thickness: -1 for filled circles
|
||||
conf_threshold: minimum confidence to draw keypoint
|
||||
Returns:
|
||||
img with keypoints drawn
|
||||
"""
|
||||
import cv2
|
||||
img = img.copy()
|
||||
for j in range(keypoints.shape[0]):
|
||||
x, y, conf = keypoints[j]
|
||||
if conf > conf_threshold:
|
||||
# Draw circle at keypoint location
|
||||
cv2.circle(img, (int(x), int(y)), radius, color, thickness)
|
||||
return img
|
||||
|
||||
def render(
|
||||
self, render_bg=True, render_ground=True, render_cam=True, fac=1.0, render_keypoints=False, **kwargs
|
||||
):
|
||||
flags = RenderFlags.RGBA | RenderFlags.SHADOWS_DIRECTIONAL
|
||||
self.ground_node.mesh.is_visible = render_ground and (not render_bg)
|
||||
if self.cam_marker_node is not None:
|
||||
self.cam_marker_node.mesh.is_visible = render_cam
|
||||
|
||||
rgba = np.zeros((*self.viewport_size[::-1], 4))
|
||||
if self.check_mesh_visibility(): # if any meshes are visible
|
||||
rgba, _ = self.viewer.render(self.scene, flags=flags)
|
||||
|
||||
img = rgba
|
||||
if render_bg and len(self.bg_seq) > 0:
|
||||
t = min(self.anim_idx, len(self.bg_seq) - 1)
|
||||
bg = self.bg_seq[t]
|
||||
rgba = rgba.astype(np.float32) / 255
|
||||
alpha = fac * rgba[..., 3:]
|
||||
img = alpha * rgba[..., :3] + (1 - alpha) * bg
|
||||
img = (255 * img).astype(np.uint8)
|
||||
|
||||
# Overlay 2D keypoints if available and requested
|
||||
if render_keypoints and render_bg:
|
||||
# Draw GT keypoints in GREEN (larger)
|
||||
if self.keypoints_seq is not None:
|
||||
t = min(self.anim_idx, len(self.keypoints_seq) - 1)
|
||||
keypoints = self.keypoints_seq[t]
|
||||
if keypoints is not None and len(keypoints) > 0:
|
||||
img = self.draw_keypoints_on_image(img, keypoints, color=(0, 255, 0), radius=5)
|
||||
|
||||
# Draw predicted keypoints in RED (smaller)
|
||||
if self.pred_keypoints_seq is not None:
|
||||
t = min(self.anim_idx, len(self.pred_keypoints_seq) - 1)
|
||||
pred_kp = self.pred_keypoints_seq[t]
|
||||
if pred_kp is not None and len(pred_kp) > 0:
|
||||
# Add confidence=1.0 if only (J, 2)
|
||||
if pred_kp.ndim == 2 and pred_kp.shape[1] == 2:
|
||||
pred_kp = np.concatenate([pred_kp, np.ones((len(pred_kp), 1))], axis=1)
|
||||
img = self.draw_keypoints_on_image(img, pred_kp, color=(0, 0, 255), radius=3, conf_threshold=0.0)
|
||||
|
||||
return img
|
||||
|
||||
def render_mesh_layers(self, render_ground=True, render_cam=True, **kwargs):
|
||||
# render each mesh in its own layer
|
||||
flags = RenderFlags.RGBA | RenderFlags.SHADOWS_DIRECTIONAL
|
||||
render_cam = render_cam and self.cam_marker_node is not None
|
||||
|
||||
# current timestep
|
||||
t = self.anim_idx % self.anim_len
|
||||
meshes = self.anim_meshes[t]
|
||||
layers = []
|
||||
|
||||
# toggle ground and camera
|
||||
self.ground_node.mesh.is_visible = False
|
||||
if self.cam_marker_node is not None:
|
||||
self.cam_marker_node.mesh.is_visible = False
|
||||
|
||||
# render animation meshes for current timestep
|
||||
for i in range(len(meshes)):
|
||||
for j, node in enumerate(self.anim_nodes):
|
||||
if i == j:
|
||||
node.mesh.is_visible = True
|
||||
else:
|
||||
node.mesh.is_visible = False
|
||||
rgba, _ = self.viewer.render(self.scene, flags=flags)
|
||||
layers.append(rgba)
|
||||
|
||||
if not render_ground and not render_cam:
|
||||
return layers
|
||||
|
||||
# render the ground and camera
|
||||
for node in self.anim_nodes:
|
||||
node.mesh.is_visible = False
|
||||
|
||||
if render_cam:
|
||||
self.cam_marker_node.mesh.is_visible = render_cam
|
||||
rgba, _ = self.viewer.render(self.scene, flags=flags)
|
||||
layers.append(rgba)
|
||||
self.cam_marker_node.mesh.is_visible = False
|
||||
|
||||
if render_ground:
|
||||
self.ground_node.mesh.is_visible = render_ground
|
||||
rgba, _ = self.viewer.render(self.scene, flags=flags)
|
||||
layers.append(rgba)
|
||||
|
||||
return layers
|
||||
|
||||
def render_frames_layers(self):
|
||||
frames_layers = []
|
||||
for t in range(self.anim_len):
|
||||
self.anim_idx = t
|
||||
self.update_frame()
|
||||
layers = self.render_mesh_layers()
|
||||
frames_layers.append(layers)
|
||||
return frames_layers
|
||||
|
||||
def animate(self, save_name, save_frames=False, render_layers=False, **kwargs):
|
||||
if render_layers:
|
||||
frames_layers = self.render_frames_layers()
|
||||
os.makedirs(save_name, exist_ok=True)
|
||||
for t, layers in enumerate(frames_layers):
|
||||
for l, layer in enumerate(layers):
|
||||
path = f"{save_name}/{t:06d}_{l:03d}.png"
|
||||
imageio.imwrite(path, layer)
|
||||
print(f"saved frame layers to {save_name}")
|
||||
return save_name
|
||||
|
||||
frames = self.render_frames(**kwargs)
|
||||
if len(frames) > 1:
|
||||
if save_frames:
|
||||
os.makedirs(save_name, exist_ok=True)
|
||||
for i, frame in enumerate(frames):
|
||||
path = f"{save_name}/{i:06d}.png"
|
||||
imageio.imwrite(path, frame)
|
||||
return save_name
|
||||
save_path = f"{save_name}.{self.ext}"
|
||||
os.makedirs(save_name, exist_ok=True)
|
||||
for idx, i in enumerate(frames):
|
||||
imageio.imwrite(f"{save_name}/" + f'{str(idx).zfill(6)}.jpg', i[:, :, :3])
|
||||
imageio.mimwrite(save_path, frames, fps=self.fps)
|
||||
print("wrote video to", save_path)
|
||||
return save_path
|
||||
|
||||
save_path = f"{save_name}.png"
|
||||
imageio.imwrite(save_path, frames[0])
|
||||
print("wrote image to", save_path)
|
||||
return save_path
|
||||
|
||||
def render_frames(self, **kwargs):
|
||||
print("ANIMATION LENGTH", self.anim_len)
|
||||
frames = []
|
||||
for t in range(self.anim_len):
|
||||
self.anim_idx = t
|
||||
self.update_frame()
|
||||
img = self.render(**kwargs)
|
||||
frames.append(img)
|
||||
return frames
|
||||
|
||||
def render_layers(self, save_dir, composite=False, fac=0.4):
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
|
||||
# render all frames
|
||||
frames = self.render_frames(render_bg=False, render_ground=False)
|
||||
|
||||
# render background
|
||||
self.set_mesh_visibility(False)
|
||||
bg_img = self.render(render_bg=False, render_ground=True) / 255
|
||||
|
||||
# save frames
|
||||
for t, frame in enumerate(frames):
|
||||
imageio.imwrite(f"{save_dir}/{t:06d}.png", frame)
|
||||
imageio.imwrite(f"{save_dir}/background.png", bg_img)
|
||||
|
||||
if composite:
|
||||
rgba = composite_layers(frames, bg_img, fac=fac)
|
||||
imageio.imwrite(f"{save_dir}/composite.png", (255 * rgba).astype(np.uint8))
|
||||
|
||||
|
||||
def composite_layers(frames, bg_img, fac=0.4):
|
||||
# composite front to back
|
||||
H, W = frames[0].shape[:2]
|
||||
comp = np.zeros((H, W, 3))
|
||||
vis = np.ones((H, W, 1))
|
||||
for frame in frames:
|
||||
frame = frame / 255
|
||||
alpha = fac * frame[..., 3:]
|
||||
comp += vis * alpha * frame[..., :3]
|
||||
vis *= 1 - alpha
|
||||
|
||||
comp += vis * bg_img[..., 3:] * bg_img[..., :3]
|
||||
vis *= 1 - bg_img[..., 3:]
|
||||
return comp + vis
|
||||
rgba = np.concatenate([comp, 1 - vis], axis=-1)
|
||||
return rgba
|
||||
|
||||
|
||||
class AnimationViewer(AnimationBase):
|
||||
def __init__(self, img_size, intrins=None, fps=30, num_repeats=1):
|
||||
super().__init__(img_size, intrins=intrins)
|
||||
|
||||
self.fps = fps
|
||||
self.ext = "gif"
|
||||
|
||||
def pause_play_callback(_, viewer):
|
||||
viewer.is_paused = not viewer.is_paused
|
||||
|
||||
def step_callback(_, viewer, step):
|
||||
viewer.anim_idx = (viewer.anim_idx + step) % viewer.anim_len
|
||||
|
||||
def stop_callback(_, viewer):
|
||||
viewer.do_animate = False
|
||||
|
||||
def loop_callback(_, viewer):
|
||||
viewer.num_repeats *= -1
|
||||
|
||||
registered_keys = {
|
||||
"p": (pause_play_callback, [self]),
|
||||
".": (step_callback, [self, 1]),
|
||||
",": (step_callback, [self, -1]),
|
||||
"l": (loop_callback, [self]),
|
||||
"s": (stop_callback, [self]),
|
||||
}
|
||||
|
||||
self.viewer = pyrender.Viewer(
|
||||
self.scene,
|
||||
run_in_thread=True,
|
||||
viewport_size=self.viewport_size,
|
||||
use_raymond_lighting=True,
|
||||
registered_keys=registered_keys,
|
||||
)
|
||||
|
||||
self.num_repeats = num_repeats
|
||||
self.do_animate = False
|
||||
self.is_paused = False
|
||||
|
||||
def acquire_lock(self):
|
||||
self.viewer.render_lock.acquire()
|
||||
|
||||
def release_lock(self):
|
||||
self.viewer.render_lock.release()
|
||||
|
||||
def close(self):
|
||||
self.viewer.close_external()
|
||||
|
||||
def animate(self, save_name=None, **kwargs):
|
||||
print("===================")
|
||||
print(f"PLAYING ANIMATION {self.num_repeats} TIMES")
|
||||
print("VIEWER CONTROLS:")
|
||||
print("p\tpause/play")
|
||||
print("l\tloop/unloop")
|
||||
print(".\tstep forward")
|
||||
print(",\tstep backward")
|
||||
print("s\tstop")
|
||||
print("===================")
|
||||
|
||||
self.anim_idx = 0
|
||||
self.do_animate = True
|
||||
frame_dur = 1 / self.fps
|
||||
if save_name is not None:
|
||||
self.viewer.viewer_flags["record"] = True
|
||||
|
||||
# set up initial frame
|
||||
self.update_frame()
|
||||
|
||||
while self.do_animate:
|
||||
if self.is_paused:
|
||||
self.update_frame()
|
||||
continue
|
||||
|
||||
if (
|
||||
self.num_repeats > 0
|
||||
and self.anim_idx >= self.num_repeats * self.anim_len
|
||||
):
|
||||
break
|
||||
|
||||
if save_name is not None and self.anim_idx >= self.anim_len:
|
||||
self.viewer.viewer_flags["record"] = False
|
||||
|
||||
start_time = time.time()
|
||||
self.update_frame()
|
||||
render_time = time.time() - start_time
|
||||
|
||||
self.anim_idx += 1
|
||||
sleep_time = frame_dur - render_time
|
||||
if sleep_time > 0:
|
||||
time.sleep(sleep_time)
|
||||
|
||||
if save_name is not None:
|
||||
print(f"SAVING RECORDING TO {save_name}")
|
||||
self.viewer.save_gif(f"{save_name}.{self.ext}")
|
||||
|
||||
|
||||
def make_checkerboard(
|
||||
length=25.0,
|
||||
color0=[0.8, 0.9, 0.9],
|
||||
color1=[0.6, 0.7, 0.7],
|
||||
tile_width=0.5,
|
||||
alpha=1.0,
|
||||
up="y",
|
||||
):
|
||||
v, f, _, fc = checkerboard_geometry(length, color0, color1, tile_width, alpha, up)
|
||||
return trimesh.Trimesh(v, f, face_colors=fc, process=False)
|
||||
|
||||
|
||||
def make_pyrender_camera(img_size, intrins=None):
|
||||
if intrins is not None:
|
||||
print("USING INTRINSICS CAMERA", intrins)
|
||||
fx, fy, cx, cy = intrins
|
||||
return pyrender.IntrinsicsCamera(fx, fy, cx, cy)
|
||||
|
||||
W, H = img_size
|
||||
focal = 0.5 * (H + W)
|
||||
yfov = 2 * np.arctan(0.5 * H / focal)
|
||||
print("USING PERSPECTIVE CAMERA", H, W, focal)
|
||||
return pyrender.PerspectiveCamera(yfov=yfov, aspectRatio=W / H)
|
||||
|
||||
|
||||
def make_camera_marker(radius=0.1, height=0.2, up="y", transform=None):
|
||||
"""
|
||||
:param radius (default 0.1) radius of pyramid base, diagonal of image plane
|
||||
:param height (default 0.2) height of pyramid, focal length
|
||||
:param up (default y) camera up vector
|
||||
:param transform (default None) (4, 4) cam to world transform
|
||||
"""
|
||||
verts, faces, face_colors = camera_marker_geometry(radius, height, up)
|
||||
if transform is not None:
|
||||
assert transform.shape == (4, 4)
|
||||
verts = (
|
||||
np.einsum("ij,nj->ni", transform[:3, :3], verts) + transform[None, :3, 3]
|
||||
)
|
||||
return trimesh.Trimesh(verts, faces, face_colors=face_colors, process=False)
|
||||
|
||||
|
||||
def get_light_poses(n_lights=5, elevation=np.pi / 3, dist=12):
|
||||
# get lights in a circle around origin at elevation
|
||||
thetas = elevation * np.ones(n_lights)
|
||||
phis = 2 * np.pi * np.arange(n_lights) / n_lights
|
||||
poses = []
|
||||
trans = make_translation(torch.tensor([0, 0, dist]))
|
||||
for phi, theta in zip(phis, thetas):
|
||||
rot = make_rotation(rx=-theta, ry=phi, order="xyz")
|
||||
poses.append((rot @ trans).numpy())
|
||||
return poses
|
||||
|
||||
|
||||
def get_raymond_light_poses(up="z"):
|
||||
thetas = np.pi * np.ones(3) / 6
|
||||
phis = 2 * np.pi * np.arange(3) / 3
|
||||
|
||||
poses = []
|
||||
|
||||
for phi, theta in zip(phis, thetas):
|
||||
xp = np.sin(theta) * np.cos(phi)
|
||||
yp = np.sin(theta) * np.sin(phi)
|
||||
zp = np.cos(theta)
|
||||
|
||||
if up == "y":
|
||||
z = np.array([xp, zp, -yp])
|
||||
else:
|
||||
z = np.array([xp, yp, zp])
|
||||
z = z / np.linalg.norm(z)
|
||||
x = np.array([-z[1], z[0], 0.0])
|
||||
if np.linalg.norm(x) == 0:
|
||||
x = np.array([1.0, 0.0, 0.0])
|
||||
x = x / np.linalg.norm(x)
|
||||
y = np.cross(z, x)
|
||||
|
||||
matrix = np.eye(4)
|
||||
matrix[:3, :3] = np.c_[x, y, z]
|
||||
poses.append(matrix)
|
||||
return poses
|
||||
Reference in New Issue
Block a user