526 lines
24 KiB
Python
526 lines
24 KiB
Python
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
|