Initial snapshot: hand motion pipeline (Dyn-HaMR + dex-retargeting + SPIDER)
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
@@ -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
|
||||
Reference in New Issue
Block a user