Spaces:
Running
Running
| import visualization.Animation as Animation | |
| from visualization.InverseKinematics import BasicInverseKinematics, BasicJacobianIK, InverseKinematics | |
| from visualization.Quaternions import Quaternions | |
| import visualization.BVH_mod as BVH | |
| from visualization.remove_fs import * | |
| from utils.plot_script import plot_3d_motion | |
| from utils import paramUtil | |
| from common.skeleton import Skeleton | |
| import torch | |
| from torch import nn | |
| from visualization.utils.quat import ik_rot, between, fk, ik | |
| from tqdm import tqdm | |
| def get_grot(glb, parent, offset): | |
| root_quat = np.array([[1.0, 0.0, 0.0, 0.0]]).repeat(glb.shape[0], axis=0)[:, None] | |
| local_pos = glb[:, 1:] - glb[:, parent[1:]] | |
| norm_offset = offset[1:] / np.linalg.norm(offset[1:], axis=-1, keepdims=True) | |
| norm_lpos = local_pos / np.linalg.norm(local_pos, axis=-1, keepdims=True) | |
| grot = between(norm_offset, norm_lpos) | |
| grot = np.concatenate((root_quat, grot), axis=1) | |
| grot /= np.linalg.norm(grot, axis=-1, keepdims=True) | |
| return grot | |
| class Joint2BVHConvertor: | |
| def __init__(self): | |
| self.template = BVH.load('./visualization/data/template.bvh', need_quater=True) | |
| self.re_order = [0, 1, 4, 7, 10, 2, 5, 8, 11, 3, 6, 9, 12, 15, 13, 16, 18, 20, 14, 17, 19, 21] | |
| self.re_order_inv = [0, 1, 5, 9, 2, 6, 10, 3, 7, 11, 4, 8, 12, 14, 18, 13, 15, 19, 16, 20, 17, 21] | |
| self.end_points = [4, 8, 13, 17, 21] | |
| self.template_offset = self.template.offsets.copy() | |
| self.parents = [-1, 0, 1, 2, 3, 0, 5, 6, 7, 0, 9, 10, 11, 12, 11, 14, 15, 16, 11, 18, 19, 20] | |
| def convert(self, positions, filename, iterations=10, foot_ik=True): | |
| ''' | |
| Convert the SMPL joint positions to Mocap BVH | |
| :param positions: (N, 22, 3) | |
| :param filename: Save path for resulting BVH | |
| :param iterations: iterations for optimizing rotations, 10 is usually enough | |
| :param foot_ik: whether to enfore foot inverse kinematics, removing foot slide issue. | |
| :return: | |
| ''' | |
| positions = positions[:, self.re_order] | |
| new_anim = self.template.copy() | |
| new_anim.rotations = Quaternions.id(positions.shape[:-1]) | |
| new_anim.positions = new_anim.positions[0:1].repeat(positions.shape[0], axis=-0) | |
| new_anim.positions[:, 0] = positions[:, 0] | |
| if foot_ik: | |
| positions = remove_fs(positions, None, fid_l=(3, 4), fid_r=(7, 8), interp_length=5, | |
| force_on_floor=True) | |
| ik_solver = BasicInverseKinematics(new_anim, positions, iterations=iterations, silent=True) | |
| new_anim = ik_solver() | |
| # BVH.save(filename, new_anim, names=new_anim.names, frametime=1 / 20, order='zyx', quater=True) | |
| glb = Animation.positions_global(new_anim)[:, self.re_order_inv] | |
| if filename is not None: | |
| BVH.save(filename, new_anim, names=new_anim.names, frametime=1 / 20, order='zyx', quater=True) | |
| return new_anim, glb | |
| def convert_sgd(self, positions, filename, iterations=100, foot_ik=True): | |
| ''' | |
| Convert the SMPL joint positions to Mocap BVH | |
| :param positions: (N, 22, 3) | |
| :param filename: Save path for resulting BVH | |
| :param iterations: iterations for optimizing rotations, 10 is usually enough | |
| :param foot_ik: whether to enfore foot inverse kinematics, removing foot slide issue. | |
| :return: | |
| ''' | |
| ## Positional Foot locking ## | |
| glb = positions[:, self.re_order] | |
| if foot_ik: | |
| glb = remove_fs(glb, None, fid_l=(3, 4), fid_r=(7, 8), interp_length=2, | |
| force_on_floor=True) | |
| ## Fit BVH ## | |
| new_anim = self.template.copy() | |
| new_anim.rotations = Quaternions.id(glb.shape[:-1]) | |
| new_anim.positions = new_anim.positions[0:1].repeat(glb.shape[0], axis=-0) | |
| new_anim.positions[:, 0] = glb[:, 0] | |
| anim = new_anim.copy() | |
| rot = torch.tensor(anim.rotations.qs, dtype=torch.float) | |
| pos = torch.tensor(anim.positions[:, 0, :], dtype=torch.float) | |
| offset = torch.tensor(anim.offsets, dtype=torch.float) | |
| glb = torch.tensor(glb, dtype=torch.float) | |
| ik_solver = InverseKinematics(rot, pos, offset, anim.parents, glb) | |
| print('Fixing foot contact using IK...') | |
| for i in tqdm(range(iterations)): | |
| mse = ik_solver.step() | |
| # print(i, mse) | |
| rotations = ik_solver.rotations.detach().cpu() | |
| norm = torch.norm(rotations, dim=-1, keepdim=True) | |
| rotations /= norm | |
| anim.rotations = Quaternions(rotations.numpy()) | |
| anim.rotations[:, self.end_points] = Quaternions.id((anim.rotations.shape[0], len(self.end_points))) | |
| anim.positions[:, 0, :] = ik_solver.position.detach().cpu().numpy() | |
| if filename is not None: | |
| BVH.save(filename, anim, names=new_anim.names, frametime=1 / 20, order='zyx', quater=True) | |
| # BVH.save(filename[:-3] + 'bvh', anim, names=new_anim.names, frametime=1 / 20, order='zyx', quater=True) | |
| glb = Animation.positions_global(anim)[:, self.re_order_inv] | |
| return anim, glb | |
| if __name__ == "__main__": | |
| # file = 'batch0_sample13_repeat0_len196.npy' | |
| # file = 'batch2_sample10_repeat0_len156.npy' | |
| # file = 'batch2_sample13_repeat0_len196.npy' #line #57 new_anim.positions = lpos #new_anim.positions[0:1].repeat(positions.shape[0], axis=-0) #TODO, figure out why it's important | |
| # file = 'batch1_sample12_repeat0_len196.npy' #hard case karate | |
| # file = 'batch1_sample14_repeat0_len180.npy' | |
| # file = 'batch0_sample3_repeat0_len192.npy' | |
| # file = 'batch1_sample4_repeat0_len136.npy' | |
| # file = 'batch0_sample0_repeat0_len152.npy' | |
| # path = f'/Users/yuxuanmu/project/MaskMIT/demo/cond4_topkr0.9_ts18_tau1.0_s1009/joints/{file}' | |
| # joints = np.load(path) | |
| # converter = Joint2BVHConvertor() | |
| # new_anim = converter.convert(joints, './gen_L196.mp4', foot_ik=True) | |
| folder = '/Users/yuxuanmu/project/MaskMIT/demo/cond4_topkr0.9_ts18_tau1.0_s1009' | |
| files = os.listdir(os.path.join(folder, 'joints')) | |
| files = [f for f in files if 'repeat' in f] | |
| converter = Joint2BVHConvertor() | |
| for f in tqdm(files): | |
| joints = np.load(os.path.join(folder, 'joints', f)) | |
| converter.convert(joints, os.path.join(folder, 'ik_animations', f'ik_{f}'.replace('npy', 'mp4')), foot_ik=True) |