153 lines
6.0 KiB
Python
153 lines
6.0 KiB
Python
"""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
|