310 lines
14 KiB
Python
310 lines
14 KiB
Python
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
|