From 686bd4d233866049f993f2594f672c0501203c56 Mon Sep 17 00:00:00 2001 From: "K. S. Ernest (iFire) Lee" Date: Thu, 8 Oct 2020 16:41:55 -0700 Subject: [PATCH] Run black formatter on code using: * shopt -s globstar * black **.py --- datasets/__init__.py | 39 ++- datasets/bvh_parser.py | 424 ++++++++++++++++++++++---- datasets/bvh_writer.py | 122 +++++--- datasets/combined_motion.py | 34 ++- datasets/fbx2bvh.py | 14 +- datasets/motion_dataset.py | 32 +- datasets/preprocess.py | 34 ++- datasets/split_joint.py | 29 +- eval.py | 12 +- eval_single_pair.py | 46 +-- get_error.py | 19 +- loss_record.py | 17 +- models/IK.py | 31 +- models/Kinematics.py | 148 ++++++--- models/__init__.py | 7 +- models/architecture.py | 198 ++++++++---- models/base_model.py | 43 ++- models/enc_and_dec.py | 152 +++++++--- models/integrated.py | 63 ++-- models/skeleton.py | 156 +++++++--- models/utils.py | 70 +++-- models/vanilla_gan.py | 89 ++++-- option_parser.py | 137 +++++---- options/options.py | 47 +-- utils/Animation.py | 532 +++++++++++++++++++-------------- utils/AnimationStructure.py | 245 +++++++++------ utils/BVH.py | 355 +++++++++++++--------- utils/BVH_mod.py | 260 ++++++++++------ utils/InverseKinematics.py | 581 ++++++++++++++++++++++-------------- utils/Pivots.py | 139 ++++++--- utils/Quaternions.py | 533 ++++++++++++++++++--------------- utils/Quaternions_old.py | 551 +++++++++++++++++++--------------- utils/__init__.py | 3 +- utils/animation_2d_data.py | 53 +++- utils/animation_data.py | 231 +++++++++----- utils/load_skeleton.py | 29 +- 36 files changed, 3517 insertions(+), 1958 deletions(-) diff --git a/datasets/__init__.py b/datasets/__init__.py index 07d406a..e034e0a 100644 --- a/datasets/__init__.py +++ b/datasets/__init__.py @@ -3,17 +3,40 @@ def get_character_names(args): """ Put the name of subdirectory in retargeting/datasets/Mixamo as [[names of group A], [names of group B]] """ - characters = [['Aj', 'BigVegas', 'Kaya', 'SportyGranny'], - ['Malcolm_m', 'Remy_m', 'Maria_m', 'Jasper_m', 'Knight_m', - 'Liam_m', 'ParasiteLStarkie_m', 'Pearl_m', 'Michelle_m', 'LolaB_m', - 'Pumpkinhulk_m', 'Ortiz_m', 'Paladin_m', 'James_m', 'Joe_m', - 'Olivia_m', 'Yaku_m', 'Timmy_m', 'Racer_m', 'Abe_m']] + characters = [ + ["Aj", "BigVegas", "Kaya", "SportyGranny"], + [ + "Malcolm_m", + "Remy_m", + "Maria_m", + "Jasper_m", + "Knight_m", + "Liam_m", + "ParasiteLStarkie_m", + "Pearl_m", + "Michelle_m", + "LolaB_m", + "Pumpkinhulk_m", + "Ortiz_m", + "Paladin_m", + "James_m", + "Joe_m", + "Olivia_m", + "Yaku_m", + "Timmy_m", + "Racer_m", + "Abe_m", + ], + ] else: """ To run evaluation successfully, number of characters in both groups must be the same. Repeat is okay. """ - characters = [['BigVegas', 'BigVegas', 'BigVegas', 'BigVegas'], ['Mousey_m', 'Goblin_m', 'Mremireh_m', 'Vampire_m']] + characters = [ + ["BigVegas", "BigVegas", "BigVegas", "BigVegas"], + ["Mousey_m", "Goblin_m", "Mremireh_m", "Vampire_m"], + ] tmp = characters[1][args.eval_seq] characters[1][args.eval_seq] = characters[1][0] characters[1][0] = tmp @@ -31,14 +54,14 @@ def create_dataset(args, character_names=None): def get_test_set(): - with open('./datasets/Mixamo/test_list.txt', 'r') as file: + with open("./datasets/Mixamo/test_list.txt", "r") as file: list = file.readlines() list = [f[:-1] for f in list] return list def get_train_list(): - with open('./datasets/Mixamo/train_list.txt', 'r') as file: + with open("./datasets/Mixamo/train_list.txt", "r") as file: list = file.readlines() list = [f[:-1] for f in list] return list diff --git a/datasets/bvh_parser.py b/datasets/bvh_parser.py index fd83f34..520e375 100644 --- a/datasets/bvh_parser.py +++ b/datasets/bvh_parser.py @@ -1,5 +1,6 @@ import sys import torch + sys.path.append("utils") sys.path.append(".") import BVH_mod as BVH @@ -15,17 +16,285 @@ Please start with root joint, then left leg chain, right leg chain, head chain, left shoulder chain and right shoulder chain. See the examples below. """ -corps_name_1 = ['Pelvis', 'LeftUpLeg', 'LeftLeg', 'LeftFoot', 'LeftToeBase', 'RightUpLeg', 'RightLeg', 'RightFoot', 'RightToeBase', 'Hips', 'Spine', 'Spine1', 'Spine2', 'Neck', 'Head', 'LeftShoulder', 'LeftArm', 'LeftForeArm', 'LeftHand', 'RightShoulder', 'RightArm', 'RightForeArm', 'RightHand'] -corps_name_2 = ['Hips', 'LeftUpLeg', 'LeftLeg', 'LeftFoot', 'LeftToeBase', 'LeftToe_End', 'RightUpLeg', 'RightLeg', 'RightFoot', 'RightToeBase', 'RightToe_End', 'Spine', 'Spine1', 'Spine2', 'Neck', 'Head', 'HeadTop_End', 'LeftShoulder', 'LeftArm', 'LeftForeArm', 'LeftHand', 'RightShoulder', 'RightArm', 'RightForeArm', 'RightHand'] -corps_name_3 = ['Hips', 'LeftUpLeg', 'LeftLeg', 'LeftFoot', 'RightUpLeg', 'RightLeg', 'RightFoot', 'Spine', 'Spine1', 'Neck', 'Head', 'LeftShoulder', 'LeftArm', 'LeftForeArm', 'LeftHand', 'RightShoulder', 'RightArm', 'RightForeArm', 'RightHand'] -corps_name_boss = ['Hips', 'LeftUpLeg', 'LeftLeg', 'LeftFoot', 'LeftToeBase', 'RightUpLeg', 'RightLeg', 'RightFoot', 'RightToeBase', 'Spine', 'Spine1', 'Spine2', 'Neck', 'Neck1', 'Head', 'LeftShoulder', 'LeftArm', 'LeftForeArm', 'LeftHand', 'RightShoulder', 'RightArm', 'RightForeArm', 'RightHand'] -corps_name_boss2 = ['Hips', 'LeftUpLeg', 'LeftLeg', 'LeftFoot', 'LeftToeBase', 'Left_End', 'RightUpLeg', 'RightLeg', 'RightFoot', 'RightToeBase', 'Right_End', 'Spine', 'Spine1', 'Spine2', 'Neck', 'Neck1', 'Head', 'LeftShoulder', 'LeftArm', 'LeftForeArm', 'LeftHand', 'RightShoulder', 'RightArm', 'RightForeArm', 'RightHand'] -corps_name_cmu = ['Hips', 'LHipJoint', 'LeftUpLeg', 'LeftLeg', 'LeftFoot', 'LeftToeBase', 'RHipJoint', 'RightUpLeg', 'RightLeg', 'RightFoot', 'RightToeBase', 'LowerBack', 'Spine', 'Spine1', 'Neck', 'Neck1', 'Head', 'LeftShoulder', 'LeftArm', 'LeftForeArm', 'LeftHand', 'RightShoulder', 'RightArm', 'RightForeArm', 'RightHand'] -corps_name_monkey = ['Hips', 'LeftUpLeg', 'LeftLeg', 'LeftFoot', 'LeftToeBase', 'RightUpLeg', 'RightLeg', 'RightFoot', 'RightToeBase', 'Spine', 'Spine1', 'Neck', 'Head', 'LeftShoulder', 'LeftArm', 'LeftForeArm', 'LeftHand', 'RightShoulder', 'RightArm', 'RightForeArm', 'RightHand'] -corps_name_three_arms = ['Three_Arms_Hips', 'LeftUpLeg', 'LeftLeg', 'LeftFoot', 'LeftToeBase', 'RightUpLeg', 'RightLeg', 'RightFoot', 'RightToeBase', 'Spine', 'Spine1', 'Neck', 'Head', 'LeftShoulder', 'LeftArm', 'LeftForeArm', 'LeftHand', 'RightShoulder', 'RightArm', 'RightForeArm', 'RightHand'] -corps_name_three_arms_split = ['Three_Arms_split_Hips', 'LeftUpLeg', 'LeftLeg', 'LeftFoot', 'LeftToeBase', 'RightUpLeg', 'RightLeg', 'RightFoot', 'RightToeBase', 'Spine', 'Spine1', 'Neck', 'Head', 'LeftShoulder', 'LeftArm', 'LeftForeArm', 'LeftHand', 'LeftHand_split', 'RightShoulder', 'RightArm', 'RightForeArm', 'RightHand', 'RightHand_split'] -corps_name_Prisoner = ['HipsPrisoner', 'LeftUpLeg', 'LeftLeg', 'LeftFoot', 'LeftToeBase', 'LeftToe_End', 'RightUpLeg', 'RightLeg', 'RightFoot', 'RightToeBase', 'RightToe_End', 'Spine', 'Spine1', 'Spine2', 'Neck', 'Head', 'HeadTop_End', 'LeftShoulder', 'LeftArm', 'LeftForeArm', 'LeftHand', 'RightShoulder', 'RightArm', 'RightForeArm'] -corps_name_mixamo2_m = ['Hips', 'LeftUpLeg', 'LeftLeg', 'LeftFoot', 'LeftToeBase', 'LeftToe_End', 'RightUpLeg', 'RightLeg', 'RightFoot', 'RightToeBase', 'RightToe_End', 'Spine', 'Spine1', 'Spine1_split', 'Spine2', 'Neck', 'Head', 'HeadTop_End', 'LeftShoulder', 'LeftShoulder_split', 'LeftArm', 'LeftForeArm', 'LeftHand', 'RightShoulder', 'RightShoulder_split', 'RightArm', 'RightForeArm', 'RightHand'] +corps_name_1 = [ + "Pelvis", + "LeftUpLeg", + "LeftLeg", + "LeftFoot", + "LeftToeBase", + "RightUpLeg", + "RightLeg", + "RightFoot", + "RightToeBase", + "Hips", + "Spine", + "Spine1", + "Spine2", + "Neck", + "Head", + "LeftShoulder", + "LeftArm", + "LeftForeArm", + "LeftHand", + "RightShoulder", + "RightArm", + "RightForeArm", + "RightHand", +] +corps_name_2 = [ + "Hips", + "LeftUpLeg", + "LeftLeg", + "LeftFoot", + "LeftToeBase", + "LeftToe_End", + "RightUpLeg", + "RightLeg", + "RightFoot", + "RightToeBase", + "RightToe_End", + "Spine", + "Spine1", + "Spine2", + "Neck", + "Head", + "HeadTop_End", + "LeftShoulder", + "LeftArm", + "LeftForeArm", + "LeftHand", + "RightShoulder", + "RightArm", + "RightForeArm", + "RightHand", +] +corps_name_3 = [ + "Hips", + "LeftUpLeg", + "LeftLeg", + "LeftFoot", + "RightUpLeg", + "RightLeg", + "RightFoot", + "Spine", + "Spine1", + "Neck", + "Head", + "LeftShoulder", + "LeftArm", + "LeftForeArm", + "LeftHand", + "RightShoulder", + "RightArm", + "RightForeArm", + "RightHand", +] +corps_name_boss = [ + "Hips", + "LeftUpLeg", + "LeftLeg", + "LeftFoot", + "LeftToeBase", + "RightUpLeg", + "RightLeg", + "RightFoot", + "RightToeBase", + "Spine", + "Spine1", + "Spine2", + "Neck", + "Neck1", + "Head", + "LeftShoulder", + "LeftArm", + "LeftForeArm", + "LeftHand", + "RightShoulder", + "RightArm", + "RightForeArm", + "RightHand", +] +corps_name_boss2 = [ + "Hips", + "LeftUpLeg", + "LeftLeg", + "LeftFoot", + "LeftToeBase", + "Left_End", + "RightUpLeg", + "RightLeg", + "RightFoot", + "RightToeBase", + "Right_End", + "Spine", + "Spine1", + "Spine2", + "Neck", + "Neck1", + "Head", + "LeftShoulder", + "LeftArm", + "LeftForeArm", + "LeftHand", + "RightShoulder", + "RightArm", + "RightForeArm", + "RightHand", +] +corps_name_cmu = [ + "Hips", + "LHipJoint", + "LeftUpLeg", + "LeftLeg", + "LeftFoot", + "LeftToeBase", + "RHipJoint", + "RightUpLeg", + "RightLeg", + "RightFoot", + "RightToeBase", + "LowerBack", + "Spine", + "Spine1", + "Neck", + "Neck1", + "Head", + "LeftShoulder", + "LeftArm", + "LeftForeArm", + "LeftHand", + "RightShoulder", + "RightArm", + "RightForeArm", + "RightHand", +] +corps_name_monkey = [ + "Hips", + "LeftUpLeg", + "LeftLeg", + "LeftFoot", + "LeftToeBase", + "RightUpLeg", + "RightLeg", + "RightFoot", + "RightToeBase", + "Spine", + "Spine1", + "Neck", + "Head", + "LeftShoulder", + "LeftArm", + "LeftForeArm", + "LeftHand", + "RightShoulder", + "RightArm", + "RightForeArm", + "RightHand", +] +corps_name_three_arms = [ + "Three_Arms_Hips", + "LeftUpLeg", + "LeftLeg", + "LeftFoot", + "LeftToeBase", + "RightUpLeg", + "RightLeg", + "RightFoot", + "RightToeBase", + "Spine", + "Spine1", + "Neck", + "Head", + "LeftShoulder", + "LeftArm", + "LeftForeArm", + "LeftHand", + "RightShoulder", + "RightArm", + "RightForeArm", + "RightHand", +] +corps_name_three_arms_split = [ + "Three_Arms_split_Hips", + "LeftUpLeg", + "LeftLeg", + "LeftFoot", + "LeftToeBase", + "RightUpLeg", + "RightLeg", + "RightFoot", + "RightToeBase", + "Spine", + "Spine1", + "Neck", + "Head", + "LeftShoulder", + "LeftArm", + "LeftForeArm", + "LeftHand", + "LeftHand_split", + "RightShoulder", + "RightArm", + "RightForeArm", + "RightHand", + "RightHand_split", +] +corps_name_Prisoner = [ + "HipsPrisoner", + "LeftUpLeg", + "LeftLeg", + "LeftFoot", + "LeftToeBase", + "LeftToe_End", + "RightUpLeg", + "RightLeg", + "RightFoot", + "RightToeBase", + "RightToe_End", + "Spine", + "Spine1", + "Spine2", + "Neck", + "Head", + "HeadTop_End", + "LeftShoulder", + "LeftArm", + "LeftForeArm", + "LeftHand", + "RightShoulder", + "RightArm", + "RightForeArm", +] +corps_name_mixamo2_m = [ + "Hips", + "LeftUpLeg", + "LeftLeg", + "LeftFoot", + "LeftToeBase", + "LeftToe_End", + "RightUpLeg", + "RightLeg", + "RightFoot", + "RightToeBase", + "RightToe_End", + "Spine", + "Spine1", + "Spine1_split", + "Spine2", + "Neck", + "Head", + "HeadTop_End", + "LeftShoulder", + "LeftShoulder_split", + "LeftArm", + "LeftForeArm", + "LeftHand", + "RightShoulder", + "RightShoulder_split", + "RightArm", + "RightForeArm", + "RightHand", +] # corps_name_example = ['Root', 'LeftUpLeg', ..., 'LeftToe', 'RightUpLeg', ..., 'RightToe', 'Spine', ..., 'Head', 'LeftShoulder', ..., 'LeftHand', 'RightShoulder', ..., 'RightHand'] """ @@ -33,20 +302,54 @@ Specify five end effectors' name. Please follow the same order as in 1. """ -ee_name_1 = ['LeftToeBase', 'RightToeBase', 'Head', 'LeftHand', 'RightHand'] -ee_name_2 = ['LeftToe_End', 'RightToe_End', 'HeadTop_End', 'LeftHand', 'RightHand'] -ee_name_3 = ['LeftFoot', 'RightFoot', 'Head', 'LeftHand', 'RightHand'] -ee_name_cmu = ['LeftToeBase', 'RightToeBase', 'Head', 'LeftHand', 'RightHand'] -ee_name_monkey = ['LeftToeBase', 'RightToeBase', 'Head', 'LeftHand', 'RightHand'] -ee_name_three_arms_split = ['LeftToeBase', 'RightToeBase', 'Head', 'LeftHand_split', 'RightHand_split'] -ee_name_Prisoner = ['LeftToe_End', 'RightToe_End', 'HeadTop_End', 'LeftHand', 'RightForeArm'] +ee_name_1 = ["LeftToeBase", "RightToeBase", "Head", "LeftHand", "RightHand"] +ee_name_2 = ["LeftToe_End", "RightToe_End", "HeadTop_End", "LeftHand", "RightHand"] +ee_name_3 = ["LeftFoot", "RightFoot", "Head", "LeftHand", "RightHand"] +ee_name_cmu = ["LeftToeBase", "RightToeBase", "Head", "LeftHand", "RightHand"] +ee_name_monkey = ["LeftToeBase", "RightToeBase", "Head", "LeftHand", "RightHand"] +ee_name_three_arms_split = [ + "LeftToeBase", + "RightToeBase", + "Head", + "LeftHand_split", + "RightHand_split", +] +ee_name_Prisoner = [ + "LeftToe_End", + "RightToe_End", + "HeadTop_End", + "LeftHand", + "RightForeArm", +] # ee_name_example = ['LeftToe', 'RightToe', 'Head', 'LeftHand', 'RightHand'] - -corps_names = [corps_name_1, corps_name_2, corps_name_3, corps_name_cmu, corps_name_monkey, corps_name_boss, - corps_name_boss, corps_name_three_arms, corps_name_three_arms_split, corps_name_Prisoner, corps_name_mixamo2_m] -ee_names = [ee_name_1, ee_name_2, ee_name_3, ee_name_cmu, ee_name_monkey, ee_name_1, ee_name_1, ee_name_1, ee_name_three_arms_split, ee_name_Prisoner, ee_name_2] +corps_names = [ + corps_name_1, + corps_name_2, + corps_name_3, + corps_name_cmu, + corps_name_monkey, + corps_name_boss, + corps_name_boss, + corps_name_three_arms, + corps_name_three_arms_split, + corps_name_Prisoner, + corps_name_mixamo2_m, +] +ee_names = [ + ee_name_1, + ee_name_2, + ee_name_3, + ee_name_cmu, + ee_name_monkey, + ee_name_1, + ee_name_1, + ee_name_1, + ee_name_three_arms_split, + ee_name_Prisoner, + ee_name_2, +] """ 3. Add previously added corps_name and ee_name at the end of the two above lists. @@ -70,8 +373,8 @@ def __init__(self, file_path=None, args=None, dataset=None, new_root=None): self.ee_length = [] for i, name in enumerate(self._names): - if ':' in name: - name = name[name.find(':') + 1:] + if ":" in name: + name = name[name.find(":") + 1 :] self._names[i] = name full_fill = [1] * len(corps_names) @@ -92,22 +395,22 @@ def __init__(self, file_path=None, args=None, dataset=None, new_root=None): if self.skeleton_type == 2 and full_fill[4]: self.skeleton_type = 4 - if 'Neck1' in self._names: + if "Neck1" in self._names: self.skeleton_type = 5 - if 'Left_End' in self._names: + if "Left_End" in self._names: self.skeleton_type = 6 - if 'Three_Arms_Hips' in self._names: + if "Three_Arms_Hips" in self._names: self.skeleton_type = 7 - if 'Three_Arms_Hips_split' in self._names: + if "Three_Arms_Hips_split" in self._names: self.skeleton_type = 8 - if 'LHipJoint' in self._names: + if "LHipJoint" in self._names: self.skeleton_type = 3 - if 'HipsPrisoner' in self._names: + if "HipsPrisoner" in self._names: self.skeleton_type = 9 - if 'Spine1_split' in self._names: + if "Spine1_split" in self._names: self.skeleton_type = 10 """ @@ -120,15 +423,17 @@ def __init__(self, file_path=None, args=None, dataset=None, new_root=None): if self.skeleton_type == -1: print(self._names) - raise Exception('Unknown skeleton') + raise Exception("Unknown skeleton") if self.skeleton_type == 0: self.set_new_root(1) self.details = [] for i, name in enumerate(self._names): - if ':' in name: name = name[name.find(':')+1:] - if name not in corps_names[self.skeleton_type]: self.details.append(i) + if ":" in name: + name = name[name.find(":") + 1 :] + if name not in corps_names[self.skeleton_type]: + self.details.append(i) self.joint_num = self.anim.shape[1] self.corps = [] self.simplified_name = [] @@ -142,9 +447,10 @@ def __init__(self, file_path=None, args=None, dataset=None, new_root=None): break if len(self.corps) != len(corps_names[self.skeleton_type]): - for i in self.corps: print(self._names[i], end=' ') - print(self.corps, self.skeleton_type, len(self.corps), sep='\n') - raise Exception('Problem in file', file_path) + for i in self.corps: + print(self._names[i], end=" ") + print(self.corps, self.skeleton_type, len(self.corps), sep="\n") + raise Exception("Problem in file", file_path) self.ee_id = [] for i in ee_names[self.skeleton_type]: @@ -169,18 +475,20 @@ def scale(self, alpha): global_position[1:, :] += (1 - alpha) * global_position[0, :] def rotate(self, theta, axis): - q = Quaternions(np.hstack((np.cos(theta/2), np.sin(theta/2) * axis))) + q = Quaternions(np.hstack((np.cos(theta / 2), np.sin(theta / 2) * axis))) position = self.anim.positions[:, 0, :].copy() rotation = self.anim.rotations[:, 0, :] position[1:, ...] -= position[0:-1, ...] - q_position = Quaternions(np.hstack((np.zeros((position.shape[0], 1)), position))) + q_position = Quaternions( + np.hstack((np.zeros((position.shape[0], 1)), position)) + ) q_rotation = Quaternions.from_euler(np.radians(rotation)) q_rotation = q * q_rotation q_position = q * q_position * (-q) self.anim.rotations[:, 0, :] = np.degrees(q_rotation.euler()) position = q_position.imaginaries for i in range(1, position.shape[0]): - position[i] += position[i-1] + position[i] += position[i - 1] self.anim.positions[:, 0, :] = position @property @@ -188,7 +496,8 @@ def topology(self): if self._topology is None: self._topology = self.anim.parents[self.corps].copy() for i in range(self._topology.shape[0]): - if i >= 1: self._topology[i] = self.simplify_map[self._topology[i]] + if i >= 1: + self._topology[i] = self.simplify_map[self._topology[i]] self._topology = tuple(self._topology) return self._topology @@ -250,7 +559,7 @@ def get_height(self): return res def write(self, file_path): - BVH.save(file_path, self.anim, self.names, self.frametime, order='xyz') + BVH.save(file_path, self.anim, self.names, self.frametime, order="xyz") def from_numpy(self, motions, frametime=None, quater=False): if frametime is not None: @@ -260,7 +569,12 @@ def from_numpy(self, motions, frametime=None, quater=False): self.anim.positions = positions[:, np.newaxis, :] if quater: rotations = motions[:, :-3].reshape(motions.shape[0], -1, 4) - norm = rotations[:, :, 0]**2 + rotations[:, :, 1]**2 + rotations[:, :, 2]**2 + rotations[:, :, 3]**2 + norm = ( + rotations[:, :, 0] ** 2 + + rotations[:, :, 1] ** 2 + + rotations[:, :, 2] ** 2 + + rotations[:, :, 3] ** 2 + ) norm = np.repeat(norm[:, :, np.newaxis], 4, axis=2) rotations /= norm rotations = Quaternions(rotations) @@ -276,10 +590,12 @@ def from_numpy(self, motions, frametime=None, quater=False): self.anim.rotations = rotations_full def get_ee_length(self): - if len(self.ee_length): return self.ee_length + if len(self.ee_length): + return self.ee_length degree = [0] * len(self.topology) for i in self.topology: - if i < 0: continue + if i < 0: + continue degree[i] += 1 for j in self.ee_id: @@ -304,17 +620,21 @@ def get_ee_length(self): def set_new_root(self, new_root): euler = torch.tensor(self.anim.rotations[:, 0, :], dtype=torch.float) - transform = ForwardKinematics.transform_from_euler(euler, 'xyz') + transform = ForwardKinematics.transform_from_euler(euler, "xyz") offset = torch.tensor(self.anim.offsets[new_root], dtype=torch.float) new_pos = torch.matmul(transform, offset) new_pos = new_pos.numpy() + self.anim.positions[:, 0, :] self.anim.offsets[0] = -self.anim.offsets[new_root] - self.anim.offsets[new_root] = np.zeros((3, )) + self.anim.offsets[new_root] = np.zeros((3,)) self.anim.positions[:, new_root, :] = new_pos - rot0 = Quaternions.from_euler(np.radians(self.anim.rotations[:, 0, :]), order='xyz') - rot1 = Quaternions.from_euler(np.radians(self.anim.rotations[:, new_root, :]), order='xyz') + rot0 = Quaternions.from_euler( + np.radians(self.anim.rotations[:, 0, :]), order="xyz" + ) + rot1 = Quaternions.from_euler( + np.radians(self.anim.rotations[:, new_root, :]), order="xyz" + ) new_rot1 = rot0 * rot1 - new_rot0 = (-rot1) + new_rot0 = -rot1 new_rot0 = np.degrees(new_rot0.euler()) new_rot1 = np.degrees(new_rot1.euler()) self.anim.rotations[:, 0, :] = new_rot0 @@ -331,7 +651,9 @@ def relabel(x): new_seq.append(x) vis[x] = 1 for y in range(len(vis)): - if not vis[y] and (self.anim.parents[x] == y or self.anim.parents[y] == x): + if not vis[y] and ( + self.anim.parents[x] == y or self.anim.parents[y] == x + ): relabel(y) new_parent[new_idx[y]] = new_idx[x] diff --git a/datasets/bvh_writer.py b/datasets/bvh_writer.py index 2c33eb7..8b35a76 100644 --- a/datasets/bvh_writer.py +++ b/datasets/bvh_writer.py @@ -1,4 +1,5 @@ import sys + sys.path.append("utils") import numpy as np from Quaternions import Quaternions @@ -6,84 +7,129 @@ # rotation with shape frame * J * 3 -def write_bvh(parent, offset, rotation, position, names, frametime, order, path, endsite=None): - file = open(path, 'w') +def write_bvh( + parent, offset, rotation, position, names, frametime, order, path, endsite=None +): + file = open(path, "w") frame = rotation.shape[0] joint_num = rotation.shape[1] order = order.upper() - file_string = 'HIERARCHY\n' + file_string = "HIERARCHY\n" def write_static(idx, prefix): nonlocal parent, offset, rotation, names, order, endsite, file_string if idx == 0: - name_label = 'ROOT ' + names[idx] - channel_label = 'CHANNELS 6 Xposition Yposition Zposition {}rotation {}rotation {}rotation'.format(*order) + name_label = "ROOT " + names[idx] + channel_label = "CHANNELS 6 Xposition Yposition Zposition {}rotation {}rotation {}rotation".format( + *order + ) else: - name_label = 'JOINT ' + names[idx] - channel_label = 'CHANNELS 3 {}rotation {}rotation {}rotation'.format(*order) - offset_label = 'OFFSET %.6f %.6f %.6f' % (offset[idx][0], offset[idx][1], offset[idx][2]) - - file_string += prefix + name_label + '\n' - file_string += prefix + '{\n' - file_string += prefix + '\t' + offset_label + '\n' - file_string += prefix + '\t' + channel_label + '\n' + name_label = "JOINT " + names[idx] + channel_label = "CHANNELS 3 {}rotation {}rotation {}rotation".format(*order) + offset_label = "OFFSET %.6f %.6f %.6f" % ( + offset[idx][0], + offset[idx][1], + offset[idx][2], + ) + + file_string += prefix + name_label + "\n" + file_string += prefix + "{\n" + file_string += prefix + "\t" + offset_label + "\n" + file_string += prefix + "\t" + channel_label + "\n" has_child = False - for y in range(idx+1, rotation.shape[1]): + for y in range(idx + 1, rotation.shape[1]): if parent[y] == idx: has_child = True - write_static(y, prefix + '\t') + write_static(y, prefix + "\t") if not has_child: - file_string += prefix + '\t' + 'End Site\n' - file_string += prefix + '\t' + '{\n' - file_string += prefix + '\t\t' + 'OFFSET 0 0 0\n' - file_string += prefix + '\t' + '}\n' + file_string += prefix + "\t" + "End Site\n" + file_string += prefix + "\t" + "{\n" + file_string += prefix + "\t\t" + "OFFSET 0 0 0\n" + file_string += prefix + "\t" + "}\n" - file_string += prefix + '}\n' + file_string += prefix + "}\n" - write_static(0, '') + write_static(0, "") - file_string += 'MOTION\n' + 'Frames: {}\n'.format(frame) + 'Frame Time: %.8f\n' % frametime + file_string += ( + "MOTION\n" + "Frames: {}\n".format(frame) + "Frame Time: %.8f\n" % frametime + ) for i in range(frame): - file_string += '%.6f %.6f %.6f ' % (position[i][0], position[i][1], position[i][2]) + file_string += "%.6f %.6f %.6f " % ( + position[i][0], + position[i][1], + position[i][2], + ) for j in range(joint_num): - file_string += '%.6f %.6f %.6f ' % (rotation[i][j][0], rotation[i][j][1], rotation[i][j][2]) - file_string += '\n' + file_string += "%.6f %.6f %.6f " % ( + rotation[i][j][0], + rotation[i][j][1], + rotation[i][j][2], + ) + file_string += "\n" file.write(file_string) return file_string -class BVH_writer(): +class BVH_writer: def __init__(self, edges, names): - self.parent, self.offset, self.names, self.edge2joint = build_joint_topology(edges, names) + self.parent, self.offset, self.names, self.edge2joint = build_joint_topology( + edges, names + ) self.joint_num = len(self.parent) # position, rotation with shape T * J * (3/4) - def write(self, rotations, positions, order, path, frametime=1.0/30, offset=None, root_y=None): - if order == 'quaternion': - norm = rotations[:, :, 0] ** 2 + rotations[:, :, 1] ** 2 + rotations[:, :, 2] ** 2 + rotations[:, :, 3] ** 2 + def write( + self, + rotations, + positions, + order, + path, + frametime=1.0 / 30, + offset=None, + root_y=None, + ): + if order == "quaternion": + norm = ( + rotations[:, :, 0] ** 2 + + rotations[:, :, 1] ** 2 + + rotations[:, :, 2] ** 2 + + rotations[:, :, 3] ** 2 + ) norm = np.repeat(norm[:, :, np.newaxis], 4, axis=2) rotations /= norm rotations = Quaternions(rotations) rotations = np.degrees(rotations.euler()) - order = 'xyz' + order = "xyz" rotations_full = np.zeros((rotations.shape[0], self.joint_num, 3)) for idx, edge in enumerate(self.edge2joint): if edge != -1: rotations_full[:, idx, :] = rotations[:, edge, :] - if root_y is not None: rotations_full[0, 0, 1] = root_y - - if offset is None: offset = self.offset - return write_bvh(self.parent, offset, rotations_full, positions, self.names, frametime, order, path) - - def write_raw(self, motion, order, path, frametime=1.0/30, root_y=None): + if root_y is not None: + rotations_full[0, 0, 1] = root_y + + if offset is None: + offset = self.offset + return write_bvh( + self.parent, + offset, + rotations_full, + positions, + self.names, + frametime, + order, + path, + ) + + def write_raw(self, motion, order, path, frametime=1.0 / 30, root_y=None): motion = motion.permute(1, 0).detach().cpu().numpy() positions = motion[:, -3:] rotations = motion[:, :-3] - if order == 'quaternion': + if order == "quaternion": rotations = rotations.reshape((motion.shape[0], -1, 4)) else: rotations = rotations.reshape((motion.shape[0], -1, 3)) diff --git a/datasets/combined_motion.py b/datasets/combined_motion.py index 4a35eb8..49a5bf0 100644 --- a/datasets/combined_motion.py +++ b/datasets/combined_motion.py @@ -13,6 +13,7 @@ class MixedData0(Dataset): """ Mixed data for many skeletons but one topologies """ + def __init__(self, args, motions, skeleton_idx): super(MixedData0, self).__init__() @@ -36,8 +37,11 @@ class MixedData(Dataset): """ data_gruop_num * 2 * samples """ + def __init__(self, args, datasets_groups): - device = torch.device(args.cuda_device if (torch.cuda.is_available()) else 'cpu') + device = torch.device( + args.cuda_device if (torch.cuda.is_available()) else "cpu" + ) self.final_data = [] self.length = 0 self.offsets = [] @@ -62,8 +66,8 @@ def __init__(self, args, datasets_groups): tmp.append(MotionData(new_args)) - mean = np.load('./datasets/Mixamo/mean_var/{}_mean.npy'.format(dataset)) - var = np.load('./datasets/Mixamo/mean_var/{}_var.npy'.format(dataset)) + mean = np.load("./datasets/Mixamo/mean_var/{}_mean.npy".format(dataset)) + var = np.load("./datasets/Mixamo/mean_var/{}_var.npy".format(dataset)) mean = torch.tensor(mean) mean = mean.reshape((1,) + mean.shape) var = torch.tensor(var) @@ -146,12 +150,14 @@ def __init__(self, args, characters): new_offset = torch.tensor(new_offset, dtype=torch.float) new_offset = new_offset.reshape((1,) + new_offset.shape) offsets_group.append(new_offset) - mean = np.load('./datasets/Mixamo/mean_var/{}_mean.npy'.format(character)) - var = np.load('./datasets/Mixamo/mean_var/{}_var.npy'.format(character)) + mean = np.load( + "./datasets/Mixamo/mean_var/{}_mean.npy".format(character) + ) + var = np.load("./datasets/Mixamo/mean_var/{}_var.npy".format(character)) mean = torch.tensor(mean) - mean = mean.reshape((1, ) + mean.shape) + mean = mean.reshape((1,) + mean.shape) var = torch.tensor(var) - var = var.reshape((1, ) + var.shape) + var = var.reshape((1,) + var.shape) mean_group.append(mean) var_group.append(var) @@ -171,20 +177,20 @@ def __getitem__(self, item): for j in range(len(character_group)): new_motion = self.get_item(i, j, item) if new_motion is not None: - new_motion = new_motion.reshape((1, ) + new_motion.shape) + new_motion = new_motion.reshape((1,) + new_motion.shape) new_motion = (new_motion - self.mean[i][j]) / self.var[i][j] ref_shape = new_motion res_group.append(new_motion) if ref_shape is None: - print('Bad at {}'.format(item)) + print("Bad at {}".format(item)) return None for j in range(len(character_group)): if res_group[j] is None: bad_flag = 1 res_group[j] = torch.zeros_like(ref_shape) if bad_flag: - print('Bad at {}'.format(item)) + print("Bad at {}".format(item)) res_group = torch.cat(res_group, dim=0) res.append([res_group, list(range(len(character_group)))]) @@ -195,17 +201,17 @@ def __len__(self): def get_item(self, gid, pid, id): character = self.characters[gid][pid] - path = './datasets/Mixamo/{}/'.format(character) + path = "./datasets/Mixamo/{}/".format(character) if isinstance(id, int): file = path + self.file_list[id] elif isinstance(id, str): file = id else: - raise Exception('Wrong input file type') + raise Exception("Wrong input file type") if not os.path.exists(file): - raise Exception('Cannot find file') + raise Exception("Cannot find file") file = BVH_file(file) - motion = file.to_tensor(quater=self.args.rotation == 'quaternion') + motion = file.to_tensor(quater=self.args.rotation == "quaternion") motion = motion[:, ::2] length = motion.shape[-1] length = length // 4 * 4 diff --git a/datasets/fbx2bvh.py b/datasets/fbx2bvh.py index 84a682b..fc810a6 100644 --- a/datasets/fbx2bvh.py +++ b/datasets/fbx2bvh.py @@ -4,11 +4,12 @@ import bpy import numpy as np import sys + sys.path.append(".") from os import listdir -data_path = './Mixamo/' +data_path = "./Mixamo/" directories = sorted([f for f in listdir(data_path) if not f.startswith(".")]) for d in directories: @@ -16,7 +17,7 @@ for f in files: sourcepath = data_path + d + "/" + f - dumppath = data_path+d + "/" + f.split(".fbx")[0] + ".bvh" + dumppath = data_path + d + "/" + f.split(".fbx")[0] + ".bvh" bpy.ops.import_scene.fbx(filepath=sourcepath) @@ -29,9 +30,12 @@ frame_start = action.frame_range[0] frame_end = np.max([60, frame_end]) - bpy.ops.export_anim.bvh(filepath=dumppath, - frame_start=frame_start, - frame_end=frame_end, root_transform_only=True) + bpy.ops.export_anim.bvh( + filepath=dumppath, + frame_start=frame_start, + frame_end=frame_end, + root_transform_only=True, + ) bpy.data.actions.remove(bpy.data.actions[-1]) print(data_path + d + "/" + f + " processed.") diff --git a/datasets/motion_dataset.py b/datasets/motion_dataset.py index b906c9a..8d7eac5 100644 --- a/datasets/motion_dataset.py +++ b/datasets/motion_dataset.py @@ -1,6 +1,7 @@ from torch.utils.data import Dataset import os import sys + sys.path.append("utils") import numpy as np import torch @@ -13,15 +14,16 @@ class MotionData(Dataset): Clip long dataset into fixed length window for batched training each data is a 2d tensor with shape (Joint_num*3) * Time """ + def __init__(self, args): super(MotionData, self).__init__() name = args.dataset - file_path = './datasets/Mixamo/{}.npy'.format(name) + file_path = "./datasets/Mixamo/{}.npy".format(name) if args.debug: - file_path = file_path[:-4] + '_debug' + file_path[-4:] + file_path = file_path[:-4] + "_debug" + file_path[-4:] - print('load from file {}'.format(file_path)) + print("load from file {}".format(file_path)) self.total_frame = 0 self.std_bvh = get_std_bvh(args) self.args = args @@ -37,7 +39,7 @@ def __init__(self, args): if args.normalization == 1: self.mean = torch.mean(self.data, (0, 2), keepdim=True) self.var = torch.var(self.data, (0, 2), keepdim=True) - self.var = self.var ** (1/2) + self.var = self.var ** (1 / 2) idx = self.var < 1e-5 self.var[idx] = 1 self.data = (self.data - self.mean) / self.var @@ -53,7 +55,11 @@ def __init__(self, args): self.reset_length_flag = 0 self.virtual_length = 0 - print('Window count: {}, total frame (without downsampling): {}'.format(len(self), self.total_frame)) + print( + "Window count: {}, total frame (without downsampling): {}".format( + len(self), self.total_frame + ) + ) def reset_length(self, length): self.reset_length_flag = 1 @@ -66,7 +72,8 @@ def __len__(self): return self.data.shape[0] def __getitem__(self, item): - if isinstance(item, int): item %= self.data.shape[0] + if isinstance(item, int): + item %= self.data.shape[0] if self.args.data_augment == 0 or np.random.randint(0, 2) == 0: return self.data[item] else: @@ -87,14 +94,18 @@ def get_windows(self, motions): end = begin + window_size new = motion[begin:end, :] - if self.args.rotation == 'quaternion': + if self.args.rotation == "quaternion": new = new.reshape(new.shape[0], -1, 3) rotations = new[:, :-1, :] rotations = Quaternions.from_euler(np.radians(rotations)).qs rotations = rotations.reshape(rotations.shape[0], -1) positions = new[:, -1, :] - positions = np.concatenate((new, np.zeros((new.shape[0], new.shape[1], 1))), axis=2) - new = np.concatenate((rotations, new[:, -1, :].reshape(new.shape[0], -1)), axis=1) + positions = np.concatenate( + (new, np.zeros((new.shape[0], new.shape[1], 1))), axis=2 + ) + new = np.concatenate( + (rotations, new[:, -1, :].reshape(new.shape[0], -1)), axis=1 + ) new = new[np.newaxis, ...] @@ -112,5 +123,6 @@ def denormalize(self, motion): self.var = self.var.to(motion.device) self.mean = self.mean.to(motion.device) ans = motion * self.var + self.mean - else: ans = motion + else: + ans = motion return ans diff --git a/datasets/preprocess.py b/datasets/preprocess.py index 5685b8c..5c3be78 100644 --- a/datasets/preprocess.py +++ b/datasets/preprocess.py @@ -9,20 +9,20 @@ def collect_bvh(data_path, character, files): - print('begin {}'.format(character)) + print("begin {}".format(character)) motions = [] for i, motion in enumerate(files): - if not os.path.exists(data_path + character + '/' + motion): + if not os.path.exists(data_path + character + "/" + motion): continue - file = BVH_file(data_path + character + '/' + motion) + file = BVH_file(data_path + character + "/" + motion) new_motion = file.to_tensor().permute((1, 0)).numpy() motions.append(new_motion) - save_file = data_path + character + '.npy' + save_file = data_path + character + ".npy" np.save(save_file, motions) - print('Npy file saved at {}'.format(save_file)) + print("Npy file saved at {}".format(save_file)) def write_statistics(character, path): @@ -38,27 +38,29 @@ def write_statistics(character, path): mean = mean.cpu().numpy()[0, ...] var = var.cpu().numpy()[0, ...] - np.save(path + '{}_mean.npy'.format(character), mean) - np.save(path + '{}_var.npy'.format(character), var) + np.save(path + "{}_mean.npy".format(character), mean) + np.save(path + "{}_var.npy".format(character), var) def copy_std_bvh(data_path, character, files): """ copy an arbitrary bvh file as a static information (skeleton's offset) reference """ - src = data_path + character + '/' + files[0] - dst = './datasets/Mixamo/std_bvhs/{}.bvh'.format(character) + src = data_path + character + "/" + files[0] + dst = "./datasets/Mixamo/std_bvhs/{}.bvh".format(character) copyfile(src, dst) -if __name__ == '__main__': - prefix = './datasets/Mixamo/' +if __name__ == "__main__": + prefix = "./datasets/Mixamo/" characters = [f for f in os.listdir(prefix) if os.path.isdir(pjoin(prefix, f))] - if 'std_bvhs' in characters: characters.remove('std_bvhs') - if 'mean_var' in characters: characters.remove('mean_var') + if "std_bvhs" in characters: + characters.remove("std_bvhs") + if "mean_var" in characters: + characters.remove("mean_var") - try_mkdir(pjoin(prefix, 'std_bvhs')) - try_mkdir(pjoin(prefix, 'mean_var')) + try_mkdir(pjoin(prefix, "std_bvhs")) + try_mkdir(pjoin(prefix, "mean_var")) for character in characters: data_path = pjoin(prefix, character) @@ -66,4 +68,4 @@ def copy_std_bvh(data_path, character, files): collect_bvh(prefix, character, files) copy_std_bvh(prefix, character, files) - write_statistics(character, './datasets/Mixamo/mean_var/') + write_statistics(character, "./datasets/Mixamo/mean_var/") diff --git a/datasets/split_joint.py b/datasets/split_joint.py index b6213ec..5c10d6f 100644 --- a/datasets/split_joint.py +++ b/datasets/split_joint.py @@ -7,6 +7,7 @@ import sys import os + sys.path.append(".") from option_parser import try_mkdir import numpy as np @@ -20,15 +21,15 @@ def split_joint(file_name, save_file=None): if save_file is None: save_file = file_name - target_joints = ['Spine1', 'LeftShoulder', 'RightShoulder'] + target_joints = ["Spine1", "LeftShoulder", "RightShoulder"] target_idx = [-1] * len(target_joints) anim, names, ftime = BVH.load(file_name) n_joint = len(anim.parents) for i, name in enumerate(names): - if ':' in name: - name = name[name.find(':') + 1:] + if ":" in name: + name = name[name.find(":") + 1 :] names[i] = name for j, joint in enumerate(target_joints): @@ -48,7 +49,8 @@ def split_joint(file_name, save_file=None): target_idx.append(-1) for i in range(n_joint): new_id[i] = i + bias - if i == target_idx[bias]: bias += 1 + if i == target_idx[bias]: + bias += 1 identity = np.zeros_like(anim.rotations) identity = identity[:, :1, :] @@ -63,7 +65,7 @@ def split_joint(file_name, save_file=None): new_anim.offsets.append(anim.offsets[i] / 2) new_anim.parents.append(i + bias) - new_names.append(names[i] + '_split') + new_names.append(names[i] + "_split") new_anim.offsets.append(anim.offsets[i] / 2) new_anim.rotations.append(identity) @@ -76,33 +78,34 @@ def split_joint(file_name, save_file=None): new_anim.offsets = np.array(new_anim.offsets) offset_spine = anim.offsets[target_idx[0]] + anim.offsets[target_idx[0] + 1] - new_anim.offsets[target_idx[0]:target_idx[0]+3, :] = offset_spine / 3 + new_anim.offsets[target_idx[0] : target_idx[0] + 3, :] = offset_spine / 3 new_anim.rotations = np.concatenate(new_anim.rotations, axis=1) try_mkdir(os.path.split(save_file)[0]) - BVH.save(save_file, new_anim, names=new_names, frametime=ftime, order='xyz') + BVH.save(save_file, new_anim, names=new_names, frametime=ftime, order="xyz") def batch_split(source, dest): - files = [f for f in os.listdir(source) if f.endswith('.bvh')] + files = [f for f in os.listdir(source) if f.endswith(".bvh")] try: bvh_file = BVH_file(os.path.join(source, files[0])) - if bvh_file.skeleton_type != 1: return + if bvh_file.skeleton_type != 1: + return except: return print("Working on {}".format(os.path.split(source)[-1])) try_mkdir(dest) - files = [f for f in os.listdir(source) if f.endswith('.bvh')] + files = [f for f in os.listdir(source) if f.endswith(".bvh")] for i, file in tqdm(enumerate(files), total=len(files)): in_file = os.path.join(source, file) out_file = os.path.join(dest, file) split_joint(in_file, out_file) -if __name__ == '__main__': - prefix = './datasets/Mixamo/' +if __name__ == "__main__": + prefix = "./datasets/Mixamo/" names = [f for f in os.listdir(prefix) if os.path.isdir(os.path.join(prefix, f))] for name in names: - batch_split(os.path.join(prefix, name), os.path.join(prefix, name + '_m')) + batch_split(os.path.join(prefix, name), os.path.join(prefix, name + "_m")) diff --git a/eval.py b/eval.py index 2f5180c..60d5988 100644 --- a/eval.py +++ b/eval.py @@ -7,15 +7,15 @@ from posixpath import join as pjoin -def eval(eval_seq, save_dir, test_device='cpu'): - para_path = pjoin(save_dir, 'para.txt') - with open(para_path, 'r') as para_file: +def eval(eval_seq, save_dir, test_device="cpu"): + para_path = pjoin(save_dir, "para.txt") + with open(para_path, "r") as para_file: argv_ = para_file.readline().split()[1:] args = option_parser.get_parser().parse_args(argv_) - args.cuda_device = test_device if torch.cuda.is_available() else 'cpu' + args.cuda_device = test_device if torch.cuda.is_available() else "cpu" args.is_train = False - args.rotation = 'quaternion' + args.rotation = "quaternion" args.eval_seq = eval_seq args.save_dir = save_dir character_names = get_character_names(args) @@ -30,7 +30,7 @@ def eval(eval_seq, save_dir, test_device='cpu'): model.test() -if __name__ == '__main__': +if __name__ == "__main__": parser = option_parser.get_parser() args = parser.parse_args() eval(args.eval_seq, args.save_dir, args.cuda_device) diff --git a/eval_single_pair.py b/eval_single_pair.py index 55dca7f..fe663d9 100644 --- a/eval_single_pair.py +++ b/eval_single_pair.py @@ -11,19 +11,19 @@ def eval_prepare(args): character = [] file_id = [] character_names = [] - character_names.append(args.input_bvh.split('/')[-2]) - character_names.append(args.target_bvh.split('/')[-2]) - if args.test_type == 'intra': - if character_names[0].endswith('_m'): - character = [['BigVegas', 'BigVegas'], character_names] + character_names.append(args.input_bvh.split("/")[-2]) + character_names.append(args.target_bvh.split("/")[-2]) + if args.test_type == "intra": + if character_names[0].endswith("_m"): + character = [["BigVegas", "BigVegas"], character_names] file_id = [[0, 0], [args.input_bvh, args.input_bvh]] src_id = 1 else: - character = [character_names, ['Goblin_m', 'Goblin_m']] + character = [character_names, ["Goblin_m", "Goblin_m"]] file_id = [[args.input_bvh, args.input_bvh], [0, 0]] src_id = 0 - elif args.test_type == 'cross': - if character_names[0].endswith('_m'): + elif args.test_type == "cross": + if character_names[0].endswith("_m"): character = [[character_names[1]], [character_names[0]]] file_id = [[0], [args.input_bvh]] src_id = 1 @@ -32,22 +32,22 @@ def eval_prepare(args): file_id = [[args.input_bvh], [0]] src_id = 0 else: - raise Exception('Unknown test type') + raise Exception("Unknown test type") return character, file_id, src_id def recover_space(file): - l = file.split('/') - l[-1] = l[-1].replace('_', ' ') - return '/'.join(l) + l = file.split("/") + l[-1] = l[-1].replace("_", " ") + return "/".join(l) def main(): parser = option_parser.get_parser() - parser.add_argument('--input_bvh', type=str, required=True) - parser.add_argument('--target_bvh', type=str, required=True) - parser.add_argument('--test_type', type=str, required=True) - parser.add_argument('--output_filename', type=str, required=True) + parser.add_argument("--input_bvh", type=str, required=True) + parser.add_argument("--target_bvh", type=str, required=True) + parser.add_argument("--test_type", type=str, required=True) + parser.add_argument("--output_filename", type=str, required=True) args = parser.parse_args() @@ -57,21 +57,21 @@ def main(): args.output_filename = recover_space(args.output_filename) character_names, file_id, src_id = eval_prepare(args) - input_character_name = args.input_bvh.split('/')[-2] - output_character_name = args.target_bvh.split('/')[-2] + input_character_name = args.input_bvh.split("/")[-2] + output_character_name = args.target_bvh.split("/")[-2] output_filename = args.output_filename test_device = args.cuda_device eval_seq = args.eval_seq - para_path = pjoin(args.save_dir, 'para.txt') - with open(para_path, 'r') as para_file: + para_path = pjoin(args.save_dir, "para.txt") + with open(para_path, "r") as para_file: argv_ = para_file.readline().split()[1:] args = option_parser.get_parser().parse_args(argv_) - args.cuda_device = test_device if torch.cuda.is_available() else 'cpu' + args.cuda_device = test_device if torch.cuda.is_available() else "cpu" args.is_train = False - args.rotation = 'quaternion' + args.rotation = "quaternion" args.eval_seq = eval_seq dataset = create_dataset(args, character_names) @@ -96,5 +96,5 @@ def main(): copyfile(bvh_path, output_filename) -if __name__ == '__main__': +if __name__ == "__main__": main() diff --git a/get_error.py b/get_error.py index d80ee9a..0206e50 100644 --- a/get_error.py +++ b/get_error.py @@ -1,6 +1,7 @@ import sys import os from option_parser import get_std_bvh + sys.path.append("utils") import BVH as BVH import numpy as np @@ -10,14 +11,14 @@ def full_batch(suffix, prefix): res = [] - chars = ['Mousey_m', 'Goblin_m', 'Mremireh_m', 'Vampire_m'] + chars = ["Mousey_m", "Goblin_m", "Mremireh_m", "Vampire_m"] for char in chars: res.append(batch(char, suffix, prefix)) return res def batch(char, suffix, prefix): - input_path = os.path.join(prefix, 'results/bvh') + input_path = os.path.join(prefix, "results/bvh") all_err = [] ref_file = get_std_bvh(dataset=char) @@ -28,8 +29,14 @@ def batch(char, suffix, prefix): new_p = os.path.join(input_path, char) - files = [f for f in os.listdir(new_p) if - f.endswith('_{}.bvh'.format(suffix)) and not f.endswith('_gt.bvh') and 'fix' not in f and not f.endswith('_input.bvh')] + files = [ + f + for f in os.listdir(new_p) + if f.endswith("_{}.bvh".format(suffix)) + and not f.endswith("_gt.bvh") + and "fix" not in f + and not f.endswith("_input.bvh") + ] for file in files: file_full = os.path.join(new_p, file) @@ -37,11 +44,11 @@ def batch(char, suffix, prefix): test_num += 1 index = [] for i, name in enumerate(names): - if 'virtual' in name: + if "virtual" in name: continue index.append(i) - file_ref = file_full[:-6] + '_gt.bvh' + file_ref = file_full[:-6] + "_gt.bvh" anim_ref, _, _ = BVH.load(file_ref) pos = Animation.positions_global(anim) # [T, J, 3] diff --git a/loss_record.py b/loss_record.py index 8f0c26a..2773ec3 100644 --- a/loss_record.py +++ b/loss_record.py @@ -12,23 +12,25 @@ def __init__(self, name: str, writer: SummaryWriter): self.writer = writer def add_scalar(self, val, step=None): - if step is None: step = len(self.loss_step) + if step is None: + step = len(self.loss_step) self.loss_step.append(val) self.loss_epoch_tmp.append(val) - self.writer.add_scalar('Train/step_' + self.name, val, step) + self.writer.add_scalar("Train/step_" + self.name, val, step) def epoch(self, step=None): - if step is None: step = len(self.loss_epoch) + if step is None: + step = len(self.loss_epoch) loss_avg = sum(self.loss_epoch_tmp) / len(self.loss_epoch_tmp) self.loss_epoch_tmp = [] self.loss_epoch.append(loss_avg) - self.writer.add_scalar('Train/epoch_' + self.name, loss_avg, step) + self.writer.add_scalar("Train/epoch_" + self.name, loss_avg, step) def save(self, path): loss_step = np.array(self.loss_step) loss_epoch = np.array(self.loss_epoch) - np.save(path + self.name + '_step.npy', loss_step) - np.save(path + self.name + '_epoch.npy', loss_epoch) + np.save(path + self.name + "_step.npy", loss_step) + np.save(path + self.name + "_epoch.npy", loss_epoch) class LossRecorder: @@ -37,7 +39,8 @@ def __init__(self, writer: SummaryWriter): self.writer = writer def add_scalar(self, name, val, step=None): - if isinstance(val, torch.Tensor): val = val.item() + 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) diff --git a/models/IK.py b/models/IK.py index 8685647..cca673f 100644 --- a/models/IK.py +++ b/models/IK.py @@ -4,7 +4,7 @@ from datasets.bvh_parser import BVH_file from tqdm import tqdm -sys.path.append('../utils') +sys.path.append("../utils") import BVH as BVH import Animation as Animation @@ -46,7 +46,7 @@ def get_foot_contact(file_name, ref_height): def get_ee_id_by_names(joint_names): - ees = ['RightToeBase', 'LeftToeBase', 'LeftFoot', 'RightFoot'] + ees = ["RightToeBase", "LeftToeBase", "LeftFoot", "RightFoot"] ee_id = [] for i, ee in enumerate(ees): ee_id.append(joint_names.index(ee)) @@ -76,7 +76,7 @@ def fix_foot_contact(input_file, foot_file, output_file, ref_height): while t + 1 < T and fixed[t + 1] == 1: t += 1 avg += glb[t, fidx].copy() - avg /= (t - s + 1) + avg /= t - s + 1 for j in range(s, t + 1): glb[j, fidx] = avg.copy() @@ -104,22 +104,25 @@ def fix_foot_contact(input_file, foot_file, output_file, ref_height): if not consl and not consr: continue if consl and consr: - litp = lerp(alpha(1.0 * (s - l + 1) / (L + 1)), - glb[s, fidx], glb[l, fidx]) - ritp = lerp(alpha(1.0 * (r - s + 1) / (L + 1)), - glb[s, fidx], glb[r, fidx]) - itp = lerp(alpha(1.0 * (s - l + 1) / (r - l + 1)), - ritp, litp) + litp = lerp( + alpha(1.0 * (s - l + 1) / (L + 1)), glb[s, fidx], glb[l, fidx] + ) + ritp = lerp( + alpha(1.0 * (r - s + 1) / (L + 1)), glb[s, fidx], glb[r, fidx] + ) + itp = lerp(alpha(1.0 * (s - l + 1) / (r - l + 1)), ritp, litp) glb[s, fidx] = itp.copy() continue if consl: - litp = lerp(alpha(1.0 * (s - l + 1) / (L + 1)), - glb[s, fidx], glb[l, fidx]) + litp = lerp( + alpha(1.0 * (s - l + 1) / (L + 1)), glb[s, fidx], glb[l, fidx] + ) glb[s, fidx] = litp.copy() continue if consr: - ritp = lerp(alpha(1.0 * (r - s + 1) / (L + 1)), - glb[s, fidx], glb[r, fidx]) + ritp = lerp( + alpha(1.0 * (r - s + 1) / (L + 1)), glb[s, fidx], glb[r, fidx] + ) glb[s, fidx] = ritp.copy() # glb is ready @@ -134,7 +137,7 @@ def fix_foot_contact(input_file, foot_file, output_file, ref_height): ik_solver = InverseKinematics(rot, pos, offset, anim.parents, glb) - print('Fixing foot contact using IK...') + print("Fixing foot contact using IK...") for i in tqdm(range(50)): ik_solver.step() diff --git a/models/Kinematics.py b/models/Kinematics.py index 7cdd5dd..fe5c5cf 100644 --- a/models/Kinematics.py +++ b/models/Kinematics.py @@ -14,22 +14,24 @@ def __init__(self, args, edges): self.world = args.fk_world self.pos_repr = args.pos_repr - self.quater = args.rotation == 'quaternion' + self.quater = args.rotation == "quaternion" def forward_from_raw(self, raw, offset, world=None, quater=None): - if world is None: world = self.world - if quater is None: quater = self.quater - if self.pos_repr == '3d': + if world is None: + world = self.world + if quater is None: + quater = self.quater + if self.pos_repr == "3d": position = raw[:, -3:, :] rotation = raw[:, :-3, :] - elif self.pos_repr == '4d': - raise Exception('Not support') + elif self.pos_repr == "4d": + raise Exception("Not support") if quater: rotation = rotation.reshape((rotation.shape[0], -1, 4, rotation.shape[-1])) identity = torch.tensor((1, 0, 0, 0), dtype=torch.float, device=raw.device) else: rotation = rotation.reshape((rotation.shape[0], -1, 3, rotation.shape[-1])) - identity = torch.zeros((3, ), dtype=torch.float, device=raw.device) + identity = torch.zeros((3,), dtype=torch.float, device=raw.device) identity = identity.reshape((1, 1, -1, 1)) new_shape = list(rotation.shape) new_shape[1] += 1 @@ -37,27 +39,38 @@ def forward_from_raw(self, raw, offset, world=None, quater=None): rotation_final = identity.repeat(new_shape) for i, j in enumerate(self.rotation_map): rotation_final[:, j, :, :] = rotation[:, i, :, :] - return self.forward(rotation_final, position, offset, world=world, quater=quater) + return self.forward( + rotation_final, position, offset, world=world, quater=quater + ) - ''' + """ rotation should have shape batch_size * Joint_num * (3/4) * Time position should have shape batch_size * 3 * Time offset should have shape batch_size * Joint_num * 3 output have shape batch_size * Time * Joint_num * 3 - ''' - def forward(self, rotation: torch.Tensor, position: torch.Tensor, offset: torch.Tensor, order='xyz', quater=False, world=True): - if not quater and rotation.shape[-2] != 3: raise Exception('Unexpected shape of rotation') - if quater and rotation.shape[-2] != 4: raise Exception('Unexpected shape of rotation') + """ + + def forward( + self, + rotation: torch.Tensor, + position: torch.Tensor, + offset: torch.Tensor, + order="xyz", + quater=False, + world=True, + ): + if not quater and rotation.shape[-2] != 3: + raise Exception("Unexpected shape of rotation") + if quater and rotation.shape[-2] != 4: + raise Exception("Unexpected shape of rotation") rotation = rotation.permute(0, 3, 1, 2) position = position.permute(0, 2, 1) - result = torch.empty(rotation.shape[:-1] + (3, ), device=position.device) - + result = torch.empty(rotation.shape[:-1] + (3,), device=position.device) norm = torch.norm(rotation, dim=-1, keepdim=True) - #norm[norm < 1e-10] = 1 + # norm[norm < 1e-10] = 1 rotation = rotation / norm - if quater: transform = self.transform_from_quaternion(rotation) else: @@ -71,9 +84,14 @@ def forward(self, rotation: torch.Tensor, position: torch.Tensor, offset: torch. assert i == 0 continue - transform[..., i, :, :] = torch.matmul(transform[..., pi, :, :], transform[..., i, :, :]) - result[..., i, :] = torch.matmul(transform[..., i, :, :], offset[..., i, :, :]).squeeze() - if world: result[..., i, :] += result[..., pi, :] + transform[..., i, :, :] = torch.matmul( + transform[..., pi, :, :], transform[..., i, :, :] + ) + result[..., i, :] = torch.matmul( + transform[..., i, :, :], offset[..., i, :, :] + ).squeeze() + if world: + result[..., i, :] += result[..., pi, :] return result def from_local_to_world(self, res: torch.Tensor): @@ -87,9 +105,13 @@ def from_local_to_world(self, res: torch.Tensor): @staticmethod def transform_from_euler(rotation, order): rotation = rotation / 180 * math.pi - transform = torch.matmul(ForwardKinematics.transform_from_axis(rotation[..., 1], order[1]), - ForwardKinematics.transform_from_axis(rotation[..., 2], order[2])) - transform = torch.matmul(ForwardKinematics.transform_from_axis(rotation[..., 0], order[0]), transform) + transform = torch.matmul( + ForwardKinematics.transform_from_axis(rotation[..., 1], order[1]), + ForwardKinematics.transform_from_axis(rotation[..., 2], order[2]), + ) + transform = torch.matmul( + ForwardKinematics.transform_from_axis(rotation[..., 0], order[0]), transform + ) return transform @staticmethod @@ -97,20 +119,20 @@ def transform_from_axis(euler, axis): transform = torch.empty(euler.shape[0:3] + (3, 3), device=euler.device) cos = torch.cos(euler) sin = torch.sin(euler) - cord = ord(axis) - ord('x') + cord = ord(axis) - ord("x") transform[..., cord, :] = transform[..., :, cord] = 0 transform[..., cord, cord] = 1 - if axis == 'x': + if axis == "x": transform[..., 1, 1] = transform[..., 2, 2] = cos transform[..., 1, 2] = -sin transform[..., 2, 1] = sin - if axis == 'y': + if axis == "y": transform[..., 0, 0] = transform[..., 2, 2] = cos transform[..., 0, 2] = sin transform[..., 2, 0] = -sin - if axis == 'z': + if axis == "z": transform[..., 0, 0] = transform[..., 1, 1] = cos transform[..., 0, 1] = -sin transform[..., 1, 0] = sin @@ -152,7 +174,14 @@ def transform_from_quaternion(quater: torch.Tensor): class InverseKinematics: - def __init__(self, rotations: torch.Tensor, positions: torch.Tensor, offset, parents, constrains): + def __init__( + self, + rotations: torch.Tensor, + positions: torch.Tensor, + offset, + parents, + constrains, + ): self.rotations = rotations self.rotations.requires_grad_(True) self.position = positions @@ -162,12 +191,21 @@ def __init__(self, rotations: torch.Tensor, positions: torch.Tensor, offset, par self.offset = offset self.constrains = constrains - self.optimizer = torch.optim.Adam([self.position, self.rotations], lr=1e-3, betas=(0.9, 0.999)) + self.optimizer = torch.optim.Adam( + [self.position, self.rotations], lr=1e-3, betas=(0.9, 0.999) + ) self.crit = nn.MSELoss() def step(self): self.optimizer.zero_grad() - glb = self.forward(self.rotations, self.position, self.offset, order='', quater=True, world=True) + glb = self.forward( + self.rotations, + self.position, + self.offset, + order="", + quater=True, + world=True, + ) loss = self.crit(glb, self.constrains) loss.backward() self.optimizer.step() @@ -181,21 +219,28 @@ def all_loss(self): res = [self.tloss(t).detach().numpy() for t in range(self.constrains.shape[0])] return np.array(res) - ''' + """ rotation should have shape batch_size * Joint_num * (3/4) * Time position should have shape batch_size * 3 * Time offset should have shape batch_size * Joint_num * 3 output have shape batch_size * Time * Joint_num * 3 - ''' - - def forward(self, rotation: torch.Tensor, position: torch.Tensor, offset: torch.Tensor, order='xyz', quater=False, - world=True): - ''' + """ + + def forward( + self, + rotation: torch.Tensor, + position: torch.Tensor, + offset: torch.Tensor, + order="xyz", + quater=False, + world=True, + ): + """ if not quater and rotation.shape[-2] != 3: raise Exception('Unexpected shape of rotation') if quater and rotation.shape[-2] != 4: raise Exception('Unexpected shape of rotation') rotation = rotation.permute(0, 3, 1, 2) position = position.permute(0, 2, 1) - ''' + """ result = torch.empty(rotation.shape[:-1] + (3,), device=position.device) norm = torch.norm(rotation, dim=-1, keepdim=True) @@ -214,17 +259,26 @@ def forward(self, rotation: torch.Tensor, position: torch.Tensor, offset: torch. assert i == 0 continue - result[..., i, :] = torch.matmul(transform[..., pi, :, :], offset[..., i, :, :]).squeeze() - transform[..., i, :, :] = torch.matmul(transform[..., pi, :, :], transform[..., i, :, :]) - if world: result[..., i, :] += result[..., pi, :] + result[..., i, :] = torch.matmul( + transform[..., pi, :, :], offset[..., i, :, :] + ).squeeze() + transform[..., i, :, :] = torch.matmul( + transform[..., pi, :, :], transform[..., i, :, :] + ) + if world: + result[..., i, :] += result[..., pi, :] return result @staticmethod def transform_from_euler(rotation, order): rotation = rotation / 180 * math.pi - transform = torch.matmul(ForwardKinematics.transform_from_axis(rotation[..., 1], order[1]), - ForwardKinematics.transform_from_axis(rotation[..., 2], order[2])) - transform = torch.matmul(ForwardKinematics.transform_from_axis(rotation[..., 0], order[0]), transform) + transform = torch.matmul( + ForwardKinematics.transform_from_axis(rotation[..., 1], order[1]), + ForwardKinematics.transform_from_axis(rotation[..., 2], order[2]), + ) + transform = torch.matmul( + ForwardKinematics.transform_from_axis(rotation[..., 0], order[0]), transform + ) return transform @staticmethod @@ -232,20 +286,20 @@ def transform_from_axis(euler, axis): transform = torch.empty(euler.shape[0:3] + (3, 3), device=euler.device) cos = torch.cos(euler) sin = torch.sin(euler) - cord = ord(axis) - ord('x') + cord = ord(axis) - ord("x") transform[..., cord, :] = transform[..., :, cord] = 0 transform[..., cord, cord] = 1 - if axis == 'x': + if axis == "x": transform[..., 1, 1] = transform[..., 2, 2] = cos transform[..., 1, 2] = -sin transform[..., 2, 1] = sin - if axis == 'y': + if axis == "y": transform[..., 0, 0] = transform[..., 2, 2] = cos transform[..., 0, 2] = sin transform[..., 2, 0] = -sin - if axis == 'z': + if axis == "z": transform[..., 0, 0] = transform[..., 1, 1] = cos transform[..., 0, 1] = -sin transform[..., 1, 0] = sin diff --git a/models/__init__.py b/models/__init__.py index 5f54ed2..0221d3c 100644 --- a/models/__init__.py +++ b/models/__init__.py @@ -1,8 +1,9 @@ def create_model(args, character_names, dataset): - if args.model == 'mul_top_mul_ske': - args.skeleton_info = 'concat' + if args.model == "mul_top_mul_ske": + args.skeleton_info = "concat" import models.architecture + return models.architecture.GAN_model(args, character_names, dataset) else: - raise Exception('Unimplemented model') + raise Exception("Unimplemented model") diff --git a/models/architecture.py b/models/architecture.py index bbd713e..3725055 100644 --- a/models/architecture.py +++ b/models/architecture.py @@ -1,7 +1,14 @@ from models.integrated import IntegratedModel from torch import optim import torch -from models.utils import GAN_loss, ImagePool, get_ee, Criterion_EE, Eval_Criterion, Criterion_EE_2 +from models.utils import ( + GAN_loss, + ImagePool, + get_ee, + Criterion_EE, + Eval_Criterion, + Criterion_EE_2, +) from models.base_model import BaseModel from option_parser import try_mkdir from posixpath import join as pjoin @@ -22,15 +29,21 @@ def __init__(self, args, character_names, dataset): self.args = args for i in range(self.n_topology): - model = IntegratedModel(args, dataset.joint_topologies[i], None, self.device, character_names[i]) + model = IntegratedModel( + args, dataset.joint_topologies[i], None, self.device, character_names[i] + ) self.models.append(model) self.D_para += model.D_parameters() self.G_para += model.G_parameters() if self.is_train: self.fake_pools = [] - self.optimizerD = optim.Adam(self.D_para, args.learning_rate, betas=(0.9, 0.999)) - self.optimizerG = optim.Adam(self.G_para, args.learning_rate, betas=(0.9, 0.999)) + self.optimizerD = optim.Adam( + self.D_para, args.learning_rate, betas=(0.9, 0.999) + ) + self.optimizerG = optim.Adam( + self.G_para, args.learning_rate, betas=(0.9, 0.999) + ) self.optimizers = [self.optimizerD, self.optimizerG] self.criterion_rec = torch.nn.MSELoss() self.criterion_gan = GAN_loss(args.gan_mode).to(self.device) @@ -40,11 +53,12 @@ def __init__(self, args, character_names, dataset): self.fake_pools.append(ImagePool(args.pool_size)) else: import option_parser + self.err_crit = [] for i in range(self.n_topology): self.err_crit.append(Eval_Criterion(dataset.joint_topologies[i])) self.id_test = 0 - self.bvh_path = pjoin(args.save_dir, 'results/bvh') + self.bvh_path = pjoin(args.save_dir, "results/bvh") option_parser.try_mkdir(self.bvh_path) self.writer = [] @@ -54,6 +68,7 @@ def __init__(self, args, character_names, dataset): from datasets.bvh_writer import BVH_writer from datasets.bvh_parser import BVH_file import option_parser + file = BVH_file(option_parser.get_std_bvh(dataset=char)) writer_group.append(BVH_writer(file.edges, file.names)) self.writer.append(writer_group) @@ -91,7 +106,9 @@ def forward(self): self.rnd_idx = [] for i in range(self.n_topology): - self.offset_repr.append(self.models[i].static_encoder(self.dataset.offsets[i])) + self.offset_repr.append( + self.models[i].static_encoder(self.dataset.offsets[i]) + ) # reconstruct for i in range(self.n_topology): @@ -101,18 +118,32 @@ def forward(self): motion_denorm = self.dataset.denorm(i, offset_idx, motion) self.motion_denorm.append(motion_denorm) - offsets = [self.offset_repr[i][p][offset_idx] for p in range(self.args.num_layers + 1)] + offsets = [ + self.offset_repr[i][p][offset_idx] + for p in range(self.args.num_layers + 1) + ] latent, res = self.models[i].auto_encoder(motion, offsets) res_denorm = self.dataset.denorm(i, offset_idx, res) - res_pos = self.models[i].fk.forward_from_raw(res_denorm, self.dataset.offsets[i][offset_idx]) + res_pos = self.models[i].fk.forward_from_raw( + res_denorm, self.dataset.offsets[i][offset_idx] + ) self.res_pos.append(res_pos) self.latents.append(latent) self.res.append(res) self.res_denorm.append(res_denorm) - pos = self.models[i].fk.forward_from_raw(motion_denorm, self.dataset.offsets[i][offset_idx]).detach() - ee = get_ee(pos, self.dataset.joint_topologies[i], self.dataset.ee_ids[i], - velo=self.args.ee_velo, from_root=self.args.ee_from_root) + pos = ( + self.models[i] + .fk.forward_from_raw(motion_denorm, self.dataset.offsets[i][offset_idx]) + .detach() + ) + ee = get_ee( + pos, + self.dataset.joint_topologies[i], + self.dataset.ee_ids[i], + velo=self.args.ee_velo, + from_root=self.args.ee_from_root, + ) height = self.models[i].height[offset_idx] height = height.reshape((height.shape[0], 1, height.shape[1], 1)) ee /= height @@ -123,18 +154,34 @@ def forward(self): for src in range(self.n_topology): for dst in range(self.n_topology): if self.is_train: - rnd_idx = torch.randint(len(self.character_names[dst]), (self.latents[src].shape[0],)) + rnd_idx = torch.randint( + len(self.character_names[dst]), (self.latents[src].shape[0],) + ) else: rnd_idx = list(range(self.latents[0].shape[0])) self.rnd_idx.append(rnd_idx) - dst_offsets_repr = [self.offset_repr[dst][p][rnd_idx] for p in range(self.args.num_layers + 1)] - fake_res = self.models[dst].auto_encoder.dec(self.latents[src], dst_offsets_repr) - fake_latent = self.models[dst].auto_encoder.enc(fake_res, dst_offsets_repr) + dst_offsets_repr = [ + self.offset_repr[dst][p][rnd_idx] + for p in range(self.args.num_layers + 1) + ] + fake_res = self.models[dst].auto_encoder.dec( + self.latents[src], dst_offsets_repr + ) + fake_latent = self.models[dst].auto_encoder.enc( + fake_res, dst_offsets_repr + ) fake_res_denorm = self.dataset.denorm(dst, rnd_idx, fake_res) - fake_pos = self.models[dst].fk.forward_from_raw(fake_res_denorm, self.dataset.offsets[dst][rnd_idx]) - fake_ee = get_ee(fake_pos, self.dataset.joint_topologies[dst], self.dataset.ee_ids[dst], - velo=self.args.ee_velo, from_root=self.args.ee_from_root) + fake_pos = self.models[dst].fk.forward_from_raw( + fake_res_denorm, self.dataset.offsets[dst][rnd_idx] + ) + fake_ee = get_ee( + fake_pos, + self.dataset.joint_topologies[dst], + self.dataset.ee_ids[dst], + velo=self.args.ee_velo, + from_root=self.args.ee_from_root, + ) height = self.models[dst].height[rnd_idx] height = height.reshape((height.shape[0], 1, height.shape[1], 1)) fake_ee = fake_ee / height @@ -145,7 +192,6 @@ def forward(self): self.fake_ee.append(fake_ee) self.fake_res_denorm.append(fake_res_denorm) - def backward_D_basic(self, netD, real, fake): """Calculate GAN loss for the discriminator Parameters: @@ -174,12 +220,16 @@ def backward_D(self): """ for i in range(self.n_topology): fake = self.fake_pools[i].query(self.fake_pos[2 - i]) - self.loss_Ds.append(self.backward_D_basic(self.models[i].discriminator, self.pos_ref[i].detach(), fake)) + self.loss_Ds.append( + self.backward_D_basic( + self.models[i].discriminator, self.pos_ref[i].detach(), fake + ) + ) self.loss_D += self.loss_Ds[-1] - self.loss_recoder.add_scalar('D_loss_{}'.format(i), self.loss_Ds[-1]) + self.loss_recoder.add_scalar("D_loss_{}".format(i), self.loss_Ds[-1]) def backward_G(self): - #rec_loss and gan loss + # rec_loss and gan loss self.rec_losses = [] self.rec_loss = 0 self.cycle_loss = 0 @@ -188,22 +238,32 @@ def backward_G(self): self.loss_G_total = 0 for i in range(self.n_topology): rec_loss1 = self.criterion_rec(self.motions[i], self.res[i]) - self.loss_recoder.add_scalar('rec_loss_quater_{}'.format(i), rec_loss1) + self.loss_recoder.add_scalar("rec_loss_quater_{}".format(i), rec_loss1) height = self.models[i].real_height[self.motions_input[i][1]] height = height.reshape(height.shape + (1, 1,)) input_pos = self.motion_denorm[i][:, -3:, :] / height rec_pos = self.res_denorm[i][:, -3:, :] / height rec_loss2 = self.criterion_rec(input_pos, rec_pos) - self.loss_recoder.add_scalar('rec_loss_global_{}'.format(i), rec_loss2) - - pos_ref_global = self.models[i].fk.from_local_to_world(self.pos_ref[i]) / height.reshape(height.shape + (1, )) - res_pos_global = self.models[i].fk.from_local_to_world(self.res_pos[i]) / height.reshape(height.shape + (1, )) + self.loss_recoder.add_scalar("rec_loss_global_{}".format(i), rec_loss2) + + pos_ref_global = self.models[i].fk.from_local_to_world( + self.pos_ref[i] + ) / height.reshape(height.shape + (1,)) + res_pos_global = self.models[i].fk.from_local_to_world( + self.res_pos[i] + ) / height.reshape(height.shape + (1,)) rec_loss3 = self.criterion_rec(pos_ref_global, res_pos_global) - self.loss_recoder.add_scalar('rec_loss_position_{}'.format(i), rec_loss3) + self.loss_recoder.add_scalar("rec_loss_position_{}".format(i), rec_loss3) - rec_loss = rec_loss1 + (rec_loss2 * self.args.lambda_global_pose + - rec_loss3 * self.args.lambda_position) * 100 + rec_loss = ( + rec_loss1 + + ( + rec_loss2 * self.args.lambda_global_pose + + rec_loss3 * self.args.lambda_position + ) + * 100 + ) self.rec_losses.append(rec_loss) self.rec_loss += rec_loss @@ -211,28 +271,38 @@ def backward_G(self): p = 0 for src in range(self.n_topology): for dst in range(self.n_topology): - cycle_loss = self.criterion_cycle(self.latents[src], self.fake_latent[p]) - self.loss_recoder.add_scalar('cycle_loss_{}_{}'.format(src, dst), cycle_loss) + cycle_loss = self.criterion_cycle( + self.latents[src], self.fake_latent[p] + ) + self.loss_recoder.add_scalar( + "cycle_loss_{}_{}".format(src, dst), cycle_loss + ) self.cycle_loss += cycle_loss ee_loss = self.criterion_ee(self.ee_ref[src], self.fake_ee[p]) - self.loss_recoder.add_scalar('ee_loss_{}_{}'.format(src, dst), ee_loss) + self.loss_recoder.add_scalar("ee_loss_{}_{}".format(src, dst), ee_loss) self.ee_loss += ee_loss if src != dst: - if self.args.gan_mode != 'none': - loss_G = self.criterion_gan(self.models[dst].discriminator(self.fake_pos[p]), True) + if self.args.gan_mode != "none": + loss_G = self.criterion_gan( + self.models[dst].discriminator(self.fake_pos[p]), True + ) else: loss_G = torch.tensor(0) - self.loss_recoder.add_scalar('G_loss_{}_{}'.format(src, dst), loss_G) + self.loss_recoder.add_scalar( + "G_loss_{}_{}".format(src, dst), loss_G + ) self.loss_G += loss_G p += 1 - self.loss_G_total = self.rec_loss * self.args.lambda_rec + \ - self.cycle_loss * self.args.lambda_cycle / 2 + \ - self.ee_loss * self.args.lambda_ee / 2 + \ - self.loss_G * 1 + self.loss_G_total = ( + self.rec_loss * self.args.lambda_rec + + self.cycle_loss * self.args.lambda_cycle / 2 + + self.ee_loss * self.args.lambda_ee / 2 + + self.loss_G * 1 + ) self.loss_G_total.backward() def optimize_parameters(self): @@ -245,7 +315,7 @@ def optimize_parameters(self): self.optimizerG.step() # update Ds - if self.args.gan_mode != 'none': + if self.args.gan_mode != "none": self.discriminator_requires_grad_(True) self.optimizerD.zero_grad() self.backward_D() @@ -254,30 +324,38 @@ def optimize_parameters(self): self.loss_D = torch.tensor(0) def verbose(self): - res = {'rec_loss_0': self.rec_losses[0].item(), - 'rec_loss_1': self.rec_losses[1].item(), - 'cycle_loss': self.cycle_loss.item(), - 'ee_loss': self.ee_loss.item(), - 'D_loss_gan': self.loss_D.item(), - 'G_loss_gan': self.loss_G.item()} + res = { + "rec_loss_0": self.rec_losses[0].item(), + "rec_loss_1": self.rec_losses[1].item(), + "cycle_loss": self.cycle_loss.item(), + "ee_loss": self.ee_loss.item(), + "D_loss_gan": self.loss_D.item(), + "G_loss_gan": self.loss_G.item(), + } return sorted(res.items(), key=lambda x: x[0]) def save(self): for i, model in enumerate(self.models): - model.save(pjoin(self.model_save_dir, 'topology{}'.format(i)), self.epoch_cnt) + model.save( + pjoin(self.model_save_dir, "topology{}".format(i)), self.epoch_cnt + ) for i, optimizer in enumerate(self.optimizers): - file_name = pjoin(self.model_save_dir, 'optimizers/{}/{}.pt'.format(self.epoch_cnt, i)) + file_name = pjoin( + self.model_save_dir, "optimizers/{}/{}.pt".format(self.epoch_cnt, i) + ) try_mkdir(psplit(file_name)[0]) torch.save(optimizer.state_dict(), file_name) def load(self, epoch=None): for i, model in enumerate(self.models): - model.load(pjoin(self.model_save_dir, 'topology{}'.format(i)), epoch) + model.load(pjoin(self.model_save_dir, "topology{}".format(i)), epoch) if self.is_train: for i, optimizer in enumerate(self.optimizers): - file_name = pjoin(self.model_save_dir, 'optimizers/{}/{}.pt'.format(epoch, i)) + file_name = pjoin( + self.model_save_dir, "optimizers/{}/{}.pt".format(epoch, i) + ) optimizer.load_state_dict(torch.load(file_name)) self.epoch_cnt = epoch @@ -289,21 +367,31 @@ def compute_test_result(self): idx = list(range(gt.shape[0])) gt = self.dataset.denorm(src, idx, gt) gt_denorm.append(gt) - gt_pose = self.models[src].fk.forward_from_raw(gt, self.dataset.offsets[src][idx]) + gt_pose = self.models[src].fk.forward_from_raw( + gt, self.dataset.offsets[src][idx] + ) gt_poses.append(gt_pose) for i in idx: new_path = pjoin(self.bvh_path, self.character_names[src][i]) from option_parser import try_mkdir + try_mkdir(new_path) - self.writer[src][i].write_raw(gt[i, ...], 'quaternion', pjoin(new_path, '{}_gt.bvh'.format(self.id_test))) + self.writer[src][i].write_raw( + gt[i, ...], + "quaternion", + pjoin(new_path, "{}_gt.bvh".format(self.id_test)), + ) p = 0 for src in range(self.n_topology): for dst in range(self.n_topology): for i in range(len(self.character_names[dst])): dst_path = pjoin(self.bvh_path, self.character_names[dst][i]) - self.writer[dst][i].write_raw(self.fake_res_denorm[p][i, ...], 'quaternion', - pjoin(dst_path, '{}_{}.bvh'.format(self.id_test, src))) + self.writer[dst][i].write_raw( + self.fake_res_denorm[p][i, ...], + "quaternion", + pjoin(dst_path, "{}_{}.bvh".format(self.id_test, src)), + ) p += 1 self.id_test += 1 diff --git a/models/base_model.py b/models/base_model.py index f961e28..1be64fc 100644 --- a/models/base_model.py +++ b/models/base_model.py @@ -17,13 +17,18 @@ class BaseModel(ABC): def __init__(self, args): self.args = args self.is_train = args.is_train - self.device = torch.device(args.cuda_device if (torch.cuda.is_available()) else 'cpu') - self.model_save_dir = pjoin(args.save_dir, 'models') # save all the checkpoints to save_dir + self.device = torch.device( + args.cuda_device if (torch.cuda.is_available()) else "cpu" + ) + self.model_save_dir = pjoin( + args.save_dir, "models" + ) # save all the checkpoints to save_dir if self.is_train: from loss_record import LossRecorder from torch.utils.tensorboard import SummaryWriter - self.log_path = pjoin(args.save_dir, 'logs') + + self.log_path = pjoin(args.save_dir, "logs") self.writer = SummaryWriter(self.log_path) self.loss_recoder = LossRecorder(self.writer) @@ -51,25 +56,35 @@ def forward(self): """Run forward pass; called by both functions and .""" 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 == 'linear': + if self.args.scheduler == "linear": + def lambda_rule(epoch): - lr_l = 1.0 - max(0, epoch - self.args.n_epochs_origin) / float(self.args.n_epochs_decay + 1) + 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 == 'Step_LR': - print('Step_LR scheduler set') + if self.args.scheduler == "Step_LR": + print("Step_LR scheduler set") return torch.optim.lr_scheduler.StepLR(optimizer, 50, 0.5) - if self.args.scheduler == 'Plateau': - print('Plateau_LR shceduler set') - return torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.2, threshold=0.01, patience=5, verbose=True) - if self.args.scheduler == 'MultiStep': + if self.args.scheduler == "Plateau": + print("Plateau_LR shceduler set") + return torch.optim.lr_scheduler.ReduceLROnPlateau( + optimizer, + mode="min", + factor=0.2, + threshold=0.01, + patience=5, + verbose=True, + ) + if self.args.scheduler == "MultiStep": return torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[]) def setup(self): @@ -78,7 +93,9 @@ def setup(self): 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] + self.schedulers = [ + self.get_scheduler(optimizer) for optimizer in self.optimizers + ] def epoch(self): self.loss_recoder.epoch() diff --git a/models/enc_and_dec.py b/models/enc_and_dec.py index 2c9da3f..64092e6 100644 --- a/models/enc_and_dec.py +++ b/models/enc_and_dec.py @@ -1,14 +1,22 @@ import torch import torch.nn as nn -from models.skeleton import SkeletonUnpool, SkeletonPool, SkeletonConv, find_neighbor, SkeletonLinear +from models.skeleton import ( + SkeletonUnpool, + SkeletonPool, + SkeletonConv, + find_neighbor, + SkeletonLinear, +) class Encoder(nn.Module): def __init__(self, args, topology): super(Encoder, self).__init__() self.topologies = [topology] - if args.rotation == 'euler_angle': self.channel_base = [3] - elif args.rotation == 'quaternion': self.channel_base = [4] + if args.rotation == "euler_angle": + self.channel_base = [3] + elif args.rotation == "quaternion": + self.channel_base = [4] self.channel_list = [] self.edge_num = [len(topology) + 1] self.pooling_list = [] @@ -19,8 +27,10 @@ def __init__(self, args, topology): kernel_size = args.kernel_size padding = (kernel_size - 1) // 2 bias = True - if args.skeleton_info == 'concat': add_offset = True - else: add_offset = False + if args.skeleton_info == "concat": + add_offset = True + else: + add_offset = False for i in range(args.num_layers): self.channel_base.append(self.channel_base[-1] * 2) @@ -29,22 +39,48 @@ def __init__(self, args, topology): seq = [] neighbor_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) + 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) for _ in range(args.extra_conv): - seq.append(SkeletonConv(neighbor_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(SkeletonConv(neighbor_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=add_offset, - in_offset_channel=3 * self.channel_base[i] // self.channel_base[0])) + seq.append( + SkeletonConv( + neighbor_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( + SkeletonConv( + neighbor_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=add_offset, + in_offset_channel=3 * self.channel_base[i] // self.channel_base[0], + ) + ) self.convs.append(seq[-1]) last_pool = True if i == args.num_layers - 1 else False - pool = SkeletonPool(edges=self.topologies[i], pooling_mode=args.skeleton_pool, - channels_per_edge=out_channels // len(neighbor_list), last_pool=last_pool) + pool = SkeletonPool( + edges=self.topologies[i], + pooling_mode=args.skeleton_pool, + channels_per_edge=out_channels // len(neighbor_list), + last_pool=last_pool, + ) seq.append(pool) seq.append(nn.LeakyReLU(negative_slope=0.2)) self.layers.append(nn.Sequential(*seq)) @@ -57,11 +93,11 @@ def __init__(self, args, topology): def forward(self, input, offset=None): # padding the one zero row to global position, so each joint including global position has 4 channels as input - if self.args.rotation == 'quaternion' and self.args.pos_repr != '4d': + if self.args.rotation == "quaternion" and self.args.pos_repr != "4d": input = torch.cat((input, torch.zeros_like(input[:, [0], :])), dim=1) for i, layer in enumerate(self.layers): - if self.args.skeleton_info == 'concat' and offset is not None: + if self.args.skeleton_info == "concat" and offset is not None: self.convs[i].set_offset(offset[i]) input = layer(input) return input @@ -79,45 +115,79 @@ def __init__(self, args, enc: Encoder): kernel_size = args.kernel_size padding = (kernel_size - 1) // 2 - if args.skeleton_info == 'concat': add_offset = True - else: add_offset = False + if args.skeleton_info == "concat": + add_offset = True + else: + add_offset = False for i in range(args.num_layers): seq = [] in_channels = enc.channel_list[args.num_layers - i] out_channels = in_channels // 2 - neighbor_list = find_neighbor(enc.topologies[args.num_layers - i - 1], args.skeleton_dist) + neighbor_list = find_neighbor( + enc.topologies[args.num_layers - i - 1], args.skeleton_dist + ) if i != 0 and i != args.num_layers - 1: bias = False else: bias = True - self.unpools.append(SkeletonUnpool(enc.pooling_list[args.num_layers - i - 1], in_channels // len(neighbor_list))) + self.unpools.append( + SkeletonUnpool( + enc.pooling_list[args.num_layers - i - 1], + in_channels // len(neighbor_list), + ) + ) - seq.append(nn.Upsample(scale_factor=2, mode=args.upsampling, align_corners=False)) + seq.append( + nn.Upsample(scale_factor=2, mode=args.upsampling, align_corners=False) + ) seq.append(self.unpools[-1]) for _ in range(args.extra_conv): - seq.append(SkeletonConv(neighbor_list, in_channels=in_channels, out_channels=in_channels, - joint_num=enc.edge_num[args.num_layers - i - 1], kernel_size=kernel_size, - stride=1, - padding=padding, padding_mode=args.padding_mode, bias=bias)) - seq.append(SkeletonConv(neighbor_list, in_channels=in_channels, out_channels=out_channels, - joint_num=enc.edge_num[args.num_layers - i - 1], kernel_size=kernel_size, stride=1, - padding=padding, padding_mode=args.padding_mode, bias=bias, add_offset=add_offset, - in_offset_channel=3 * enc.channel_base[args.num_layers - i - 1] // enc.channel_base[0])) + seq.append( + SkeletonConv( + neighbor_list, + in_channels=in_channels, + out_channels=in_channels, + joint_num=enc.edge_num[args.num_layers - i - 1], + kernel_size=kernel_size, + stride=1, + padding=padding, + padding_mode=args.padding_mode, + bias=bias, + ) + ) + seq.append( + SkeletonConv( + neighbor_list, + in_channels=in_channels, + out_channels=out_channels, + joint_num=enc.edge_num[args.num_layers - i - 1], + kernel_size=kernel_size, + stride=1, + padding=padding, + padding_mode=args.padding_mode, + bias=bias, + add_offset=add_offset, + in_offset_channel=3 + * enc.channel_base[args.num_layers - i - 1] + // enc.channel_base[0], + ) + ) self.convs.append(seq[-1]) - if i != args.num_layers - 1: seq.append(nn.LeakyReLU(negative_slope=0.2)) + if i != args.num_layers - 1: + seq.append(nn.LeakyReLU(negative_slope=0.2)) self.layers.append(nn.Sequential(*seq)) def forward(self, input, offset=None): for i, layer in enumerate(self.layers): - if self.args.skeleton_info == 'concat': + if self.args.skeleton_info == "concat": self.convs[i].set_offset(offset[len(self.layers) - i - 1]) input = layer(input) # throw the padded rwo for global position - if self.args.rotation == 'quaternion' and self.args.pos_repr != '4d': + if self.args.rotation == "quaternion" and self.args.pos_repr != "4d": input = input[:, :-1, :] return input @@ -147,10 +217,18 @@ def __init__(self, args, edges): for i in range(args.num_layers): neighbor_list = find_neighbor(edges, args.skeleton_dist) seq = [] - seq.append(SkeletonLinear(neighbor_list, in_channels=channels * len(neighbor_list), - out_channels=channels * 2 * len(neighbor_list), extra_dim1=True)) + seq.append( + SkeletonLinear( + neighbor_list, + in_channels=channels * len(neighbor_list), + out_channels=channels * 2 * len(neighbor_list), + extra_dim1=True, + ) + ) if i < args.num_layers - 1: - pool = SkeletonPool(edges, channels_per_edge=channels*2, pooling_mode='mean') + pool = SkeletonPool( + edges, channels_per_edge=channels * 2, pooling_mode="mean" + ) seq.append(pool) edges = pool.new_edges seq.append(activation) diff --git a/models/integrated.py b/models/integrated.py index caeae9b..a17ae6c 100644 --- a/models/integrated.py +++ b/models/integrated.py @@ -11,20 +11,24 @@ class IntegratedModel: # origin_offsets should have shape num_skeleton * J * 3 - def __init__(self, args, joint_topology, origin_offsets: torch.Tensor, device, characters): + def __init__( + self, args, joint_topology, origin_offsets: torch.Tensor, device, characters + ): self.args = args self.joint_topology = joint_topology - self.edges = build_edge_topology(joint_topology, torch.zeros((len(joint_topology), 3))) + self.edges = build_edge_topology( + joint_topology, torch.zeros((len(joint_topology), 3)) + ) self.fk = ForwardKinematics(args, self.edges) - self.height = [] # for normalize ee_loss + self.height = [] # for normalize ee_loss self.real_height = [] for char in characters: if args.use_sep_ee: h = BVH_file(get_std_bvh(dataset=char)).get_ee_length() else: h = BVH_file(get_std_bvh(dataset=char)).get_height() - if args.ee_loss_fact == 'learn': + if args.ee_loss_fact == "learn": h = torch.tensor(h, dtype=torch.float) else: h = torch.tensor(h, dtype=torch.float, requires_grad=False) @@ -33,22 +37,29 @@ def __init__(self, args, joint_topology, origin_offsets: torch.Tensor, device, c self.real_height = torch.tensor(self.real_height, device=device) self.height = torch.cat(self.height, dim=0) self.height = self.height.to(device) - if not args.use_sep_ee: self.height.unsqueeze_(-1) - if args.ee_loss_fact == 'learn': self.height_para = [self.height] - else: self.height_para = [] + if not args.use_sep_ee: + self.height.unsqueeze_(-1) + if args.ee_loss_fact == "learn": + self.height_para = [self.height] + else: + self.height_para = [] if not args.simple_operator: self.auto_encoder = AE(args, topology=self.edges).to(device) self.discriminator = Discriminator(args, self.edges).to(device) self.static_encoder = StaticEncoder(args, self.edges).to(device) else: - raise Exception('Conventional operator not yet implemented') + raise Exception("Conventional operator not yet implemented") def parameters(self): return self.G_parameters() + self.D_parameters() def G_parameters(self): - return list(self.auto_encoder.parameters()) + list(self.static_encoder.parameters()) + self.height_para + return ( + list(self.auto_encoder.parameters()) + + list(self.static_encoder.parameters()) + + self.height_para + ) def D_parameters(self): return list(self.discriminator.parameters()) @@ -58,29 +69,35 @@ def save(self, path, epoch): path = pjoin(path, str(epoch)) try_mkdir(path) - torch.save(self.height, pjoin(path, 'height.pt')) - torch.save(self.auto_encoder.state_dict(), pjoin(path, 'auto_encoder.pt')) - torch.save(self.discriminator.state_dict(), pjoin(path, 'discriminator.pt')) - torch.save(self.static_encoder.state_dict(), pjoin(path, 'static_encoder.pt')) + torch.save(self.height, pjoin(path, "height.pt")) + torch.save(self.auto_encoder.state_dict(), pjoin(path, "auto_encoder.pt")) + torch.save(self.discriminator.state_dict(), pjoin(path, "discriminator.pt")) + torch.save(self.static_encoder.state_dict(), pjoin(path, "static_encoder.pt")) - print('Save at {} succeed!'.format(path)) + print("Save at {} succeed!".format(path)) def load(self, path, epoch=None): - print('loading from', path) + print("loading from", path) if not os.path.exists(path): - raise Exception('Unknown loading path') + raise Exception("Unknown loading path") if epoch is None: all = [int(q) for q in os.listdir(path) if os.path.isdir(path + q)] if len(all) == 0: - raise Exception('Empty loading path') + raise Exception("Empty loading path") epoch = sorted(all)[-1] path = pjoin(path, str(epoch)) - print('loading from epoch {}......'.format(epoch)) + print("loading from epoch {}......".format(epoch)) - self.auto_encoder.load_state_dict(torch.load(pjoin(path, 'auto_encoder.pt'), - map_location=self.args.cuda_device)) - self.static_encoder.load_state_dict(torch.load(pjoin(path, 'static_encoder.pt'), - map_location=self.args.cuda_device)) - print('load succeed!') + self.auto_encoder.load_state_dict( + torch.load( + pjoin(path, "auto_encoder.pt"), map_location=self.args.cuda_device + ) + ) + self.static_encoder.load_state_dict( + torch.load( + pjoin(path, "static_encoder.pt"), map_location=self.args.cuda_device + ) + ) + print("load succeed!") diff --git a/models/skeleton.py b/models/skeleton.py index 1a09158..cf34274 100644 --- a/models/skeleton.py +++ b/models/skeleton.py @@ -6,16 +6,30 @@ 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): + 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') + raise Exception("BAD") super(SkeletonConv, self).__init__() - if padding_mode == 'zeros': padding_mode = 'constant' - if padding_mode == 'reflection': padding_mode = 'reflect' + if padding_mode == "zeros": + padding_mode = "constant" + if padding_mode == "reflection": + padding_mode = "reflect" self.expanded_neighbour_list = [] self.expanded_neighbour_list_offset = [] @@ -38,7 +52,9 @@ def __init__(self, neighbour_list, in_channels, out_channels, kernel_size, joint self.expanded_neighbour_list.append(expanded) if self.add_offset: - self.offset_enc = SkeletonLinear(neighbour_list, in_offset_channel * len(neighbour_list), out_channels) + self.offset_enc = SkeletonLinear( + neighbour_list, in_offset_channel * len(neighbour_list), out_channels + ) for neighbour in neighbour_list: expanded = [] @@ -51,16 +67,28 @@ def __init__(self, neighbour_list, in_channels, out_channels, kernel_size, joint if bias: self.bias = torch.zeros(out_channels) else: - self.register_parameter('bias', None) + 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[ + 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.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() @@ -68,33 +96,66 @@ def __init__(self, neighbour_list, in_channels, out_channels, kernel_size, joint 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 un expected result """ - tmp = torch.zeros_like(self.weight[self.out_channels_per_joint * i: self.out_channels_per_joint * (i + 1), - neighbour, ...]) + 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 + 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, ...]) + 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)]) + 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.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') + if not self.add_offset: + raise Exception("Wrong Combination of Parameters") self.offset = offset.reshape(offset.shape[0], -1) def forward(self, input): weight_masked = self.weight * self.mask - 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) + 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, + ) return res @@ -125,11 +186,22 @@ def __init__(self, neighbour_list, in_channels, out_channels, extra_dim1=False): 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.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 + 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 + 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) @@ -142,7 +214,8 @@ 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,)) + if self.extra_dim1: + res = res.reshape(res.shape + (1,)) return res @@ -150,8 +223,8 @@ 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') + if pooling_mode != "mean": + raise Exception("Unimplemented pooling mode in matrix_implementation") self.channels_per_edge = channels_per_edge self.pooling_mode = pooling_mode @@ -196,16 +269,21 @@ def find_seq(j, seq): # add global position self.pooling_list.append([self.edge_num - 1]) - self.description = 'SkeletonPool(in_edge_num={}, out_edge_num={})'.format( + 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) + 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[ + i * channels_per_edge + c, j * channels_per_edge + c + ] = 1.0 / len(pair) self.weight = nn.Parameter(self.weight, requires_grad=False) @@ -223,16 +301,21 @@ def __init__(self, pooling_list, 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.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) + 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[ + j * channels_per_edge + c, i * channels_per_edge + c + ] = 1 self.weight = nn.Parameter(self.weight) self.weight.requires_grad_(False) @@ -245,6 +328,7 @@ def forward(self, input: torch.Tensor): Helper functions for skeleton operation """ + def dfs(x, fa, vis, dist): vis[x] = 1 for y in range(len(fa)): @@ -252,6 +336,7 @@ def dfs(x, fa, vis, dist): dist[y] = dist[x] + 1 dfs(y, fa, vis, dist) + """ def find_neighbor_joint(fa, threshold): neighbor_list = [[]] @@ -276,6 +361,7 @@ def find_neighbor_joint(fa, threshold): return neighbor_list """ + def build_edge_topology(topology, offset): # get all edges (pa, child, offset) edges = [] @@ -309,7 +395,7 @@ def make_topology(edge_idx, pa): if out_degree[edge[0]] > 1: parent.append(pa) offset.append(np.array([0, 0, 0])) - names.append(origin_names[edge[1]] + '_virtual') + names.append(origin_names[edge[1]] + "_virtual") edge2joint.append(-1) pa = joint_cnt joint_cnt += 1 diff --git a/models/utils.py b/models/utils.py index 842e5ba..1e7aa6f 100644 --- a/models/utils.py +++ b/models/utils.py @@ -10,17 +10,17 @@ class GAN_loss(nn.Module): def __init__(self, gan_mode, real_lable=1.0, fake_lable=0.0): super(GAN_loss, self).__init__() - self.register_buffer('real_label', torch.tensor(real_lable)) - self.register_buffer('fake_label', torch.tensor(fake_lable)) + self.register_buffer("real_label", torch.tensor(real_lable)) + self.register_buffer("fake_label", torch.tensor(fake_lable)) self.gan_mode = gan_mode - if gan_mode == 'lsgan': + if gan_mode == "lsgan": self.loss = nn.MSELoss() - elif gan_mode == 'vanilla': + elif gan_mode == "vanilla": self.loss = nn.BCEWithLogitsLoss() - elif gan_mode == 'none': + elif gan_mode == "none": self.loss = None else: - raise Exception('Unknown GAN mode') + raise Exception("Unknown GAN mode") def get_target_tensor(self, prediction, target_is_real): if target_is_real: @@ -54,9 +54,10 @@ def __call__(self, pred, gt): def parameters(self): return [] + class Criterion_EE_2: def __init__(self, args, base_criterion, norm_eps=0.008): - print('Using adaptive EE') + print("Using adaptive EE") self.args = args self.base_criterion = base_criterion self.norm_eps = norm_eps @@ -73,6 +74,7 @@ def __call__(self, pred, gt): def parameters(self): return list(self.ada_para.parameters()) + class Eval_Criterion: def __init__(self, parent): self.pa = parent @@ -86,7 +88,7 @@ def __call__(self, pred, gt): return self.base_criterion(pred, gt) -class ImagePool(): +class ImagePool: """This class implements an image buffer that stores previously generated images. This buffer enables us to update discriminators using a history of generated images rather than the ones produced by the latest generators. @@ -116,20 +118,26 @@ def query(self, images): return_images = [] for image in images: image = torch.unsqueeze(image.data, 0) - if self.num_imgs < self.pool_size: # if the buffer is not full; keep inserting current images to the buffer + if ( + self.num_imgs < self.pool_size + ): # if the buffer is not full; keep inserting current images to the buffer self.num_imgs = self.num_imgs + 1 self.images.append(image) return_images.append(image) else: p = random.uniform(0, 1) - if p > 0.5: # by 50% chance, the buffer will return a previously stored image, and insert the current image into the buffer - random_id = random.randint(0, self.pool_size - 1) # randint is inclusive + if ( + p > 0.5 + ): # by 50% chance, the buffer will return a previously stored image, and insert the current image into the buffer + random_id = random.randint( + 0, self.pool_size - 1 + ) # randint is inclusive tmp = self.images[random_id].clone() self.images[random_id] = image return_images.append(tmp) - else: # by another 50% chance, the buffer will return the current image + else: # by another 50% chance, the buffer will return the current image return_images.append(image) - return_images = torch.cat(return_images, 0) # collect all the images and return + return_images = torch.cat(return_images, 0) # collect all the images and return return return_images @@ -144,27 +152,41 @@ def get_scheduler(optimizer, opt): For other schedulers (step, plateau, and cosine), we use the default PyTorch schedulers. See https://pytorch.org/docs/stable/optim.html for more details. """ - if opt.lr_policy == 'linear': + if opt.lr_policy == "linear": + def lambda_rule(epoch): - lr_l = 1.0 - max(0, epoch + opt.epoch_count - opt.n_epochs) / float(opt.n_epochs_decay + 1) + lr_l = 1.0 - max(0, epoch + opt.epoch_count - opt.n_epochs) / float( + opt.n_epochs_decay + 1 + ) return lr_l + scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_rule) - elif opt.lr_policy == 'step': - scheduler = lr_scheduler.StepLR(optimizer, step_size=opt.lr_decay_iters, gamma=0.1) - elif opt.lr_policy == 'plateau': - scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.2, threshold=0.01, patience=5) - elif opt.lr_policy == 'cosine': - scheduler = lr_scheduler.CosineAnnealingLR(optimizer, T_max=opt.n_epochs, eta_min=0) + elif opt.lr_policy == "step": + scheduler = lr_scheduler.StepLR( + optimizer, step_size=opt.lr_decay_iters, gamma=0.1 + ) + elif opt.lr_policy == "plateau": + scheduler = lr_scheduler.ReduceLROnPlateau( + optimizer, mode="min", factor=0.2, threshold=0.01, patience=5 + ) + elif opt.lr_policy == "cosine": + scheduler = lr_scheduler.CosineAnnealingLR( + optimizer, T_max=opt.n_epochs, eta_min=0 + ) else: - return NotImplementedError('learning rate policy [%s] is not implemented', opt.lr_policy) + return NotImplementedError( + "learning rate policy [%s] is not implemented", opt.lr_policy + ) return scheduler def get_ee(pos, pa, ees, velo=False, from_root=False): pos = pos.clone() for i, fa in enumerate(pa): - if i == 0: continue - if not from_root and fa == 0: continue + if i == 0: + continue + if not from_root and fa == 0: + continue pos[:, :, i, :] += pos[:, :, fa, :] pos = pos[:, :, ees, :] diff --git a/models/vanilla_gan.py b/models/vanilla_gan.py index 0175b2f..814e6ed 100644 --- a/models/vanilla_gan.py +++ b/models/vanilla_gan.py @@ -25,25 +25,44 @@ def __init__(self, args, topology): for i in range(args.num_layers): seq = [] - neighbor_list = skeleton.find_neighbor(self.topologies[i], args.skeleton_dist) + neighbor_list = skeleton.find_neighbor( + self.topologies[i], args.skeleton_dist + ) in_channels = self.channel_base[i] * self.joint_num[i] - out_channels = self.channel_base[i+1] * self.joint_num[i] - if i == 0: self.channel_list.append(in_channels) + out_channels = self.channel_base[i + 1] * self.joint_num[i] + if i == 0: + self.channel_list.append(in_channels) self.channel_list.append(out_channels) - if i < args.num_layers - 1: bias = False - else: bias = True + if i < args.num_layers - 1: + bias = False + else: + bias = True if i == args.num_layers - 1: kernel_size = 16 padding = 0 - seq.append(SkeletonConv(neighbor_list, in_channels=in_channels, out_channels=out_channels, - joint_num=self.joint_num[i], kernel_size=kernel_size, stride=2, padding=padding, - padding_mode='reflection', bias=bias)) - if i < args.num_layers - 1: seq.append(nn.BatchNorm1d(out_channels)) - pool = SkeletonPool(edges=self.topologies[i], pooling_mode=args.skeleton_pool, - channels_per_edge=out_channels // len(neighbor_list)) + seq.append( + SkeletonConv( + neighbor_list, + in_channels=in_channels, + out_channels=out_channels, + joint_num=self.joint_num[i], + kernel_size=kernel_size, + stride=2, + padding=padding, + padding_mode="reflection", + bias=bias, + ) + ) + if i < args.num_layers - 1: + seq.append(nn.BatchNorm1d(out_channels)) + pool = SkeletonPool( + edges=self.topologies[i], + pooling_mode=args.skeleton_pool, + channels_per_edge=out_channels // len(neighbor_list), + ) seq.append(pool) if not self.args.patch_gan or i < args.num_layers - 1: seq.append(nn.LeakyReLU(negative_slope=0.2)) @@ -53,9 +72,10 @@ def __init__(self, args, topology): self.pooling_list.append(pool.pooling_list) self.joint_num.append(len(pool.new_edges) + 1) if i == args.num_layers - 1: - self.last_channel = self.joint_num[-1] * self.channel_base[i+1] + self.last_channel = self.joint_num[-1] * self.channel_base[i + 1] - if not args.patch_gan: self.compress = nn.Linear(in_features=self.last_channel, out_features=1) + if not args.patch_gan: + self.compress = nn.Linear(in_features=self.last_channel, out_features=1) def forward(self, input): input = input.reshape(input.shape[0], input.shape[1], -1) @@ -73,7 +93,15 @@ class Generator(nn.Module): def __init__(self, args, d: Discriminator): super(Generator, self).__init__() self.latent_dimension = args.latent_dimension - self.begin = nn.Parameter(data=torch.randn((self.latent_dimension, d.last_channel, args.window_size // (2 ** (args.num_layers - 1)))) ) + self.begin = nn.Parameter( + data=torch.randn( + ( + self.latent_dimension, + d.last_channel, + args.window_size // (2 ** (args.num_layers - 1)), + ) + ) + ) self.layers = nn.ModuleList() kernel_size = args.kernel_size padding = (kernel_size - 1) // 2 @@ -82,16 +110,35 @@ def __init__(self, args, d: Discriminator): seq = [] in_channels = d.channel_list[args.num_layers - i] out_channels = in_channels // 2 - neighbor_list = skeleton.find_neighbor(d.topologies[args.num_layers - i - 1], args.skeleton_dist) + neighbor_list = skeleton.find_neighbor( + d.topologies[args.num_layers - i - 1], args.skeleton_dist + ) - if i != 0 and i != args.num_layers - 1: bias = False - else: bias = True + if i != 0 and i != args.num_layers - 1: + bias = False + else: + bias = True - if i != 0: seq.append(nn.Upsample(scale_factor=2, mode=args.upsampling)) + if i != 0: + seq.append(nn.Upsample(scale_factor=2, mode=args.upsampling)) seq.append(SkeletonUnpool(d.pooling_maps[args.num_layers - i - 1])) - seq.append(SkeletonConv(neighbor_list, in_channels=in_channels, out_channels=out_channels, joint_num=d.joint_num[args.num_layers - i - 1], kernel_size=kernel_size, stride=1, padding=padding, padding_mode=args.padding_mode, bias=bias)) - if i != 0 and i != args.num_layers - 1: seq.append(nn.BatchNorm1d(out_channels)) - if i != args.num_layers - 1: seq.append(nn.LeakyReLU(negative_slope=0.2)) + seq.append( + SkeletonConv( + neighbor_list, + in_channels=in_channels, + out_channels=out_channels, + joint_num=d.joint_num[args.num_layers - i - 1], + kernel_size=kernel_size, + stride=1, + padding=padding, + padding_mode=args.padding_mode, + bias=bias, + ) + ) + if i != 0 and i != args.num_layers - 1: + seq.append(nn.BatchNorm1d(out_channels)) + if i != args.num_layers - 1: + seq.append(nn.LeakyReLU(negative_slope=0.2)) self.layers.append(nn.Sequential(*seq)) diff --git a/option_parser.py b/option_parser.py index b875549..ab2b7bf 100644 --- a/option_parser.py +++ b/option_parser.py @@ -3,54 +3,86 @@ def get_parser(): parser = argparse.ArgumentParser() - parser.add_argument('--save_dir', type=str, default='./pretrained', help='directory for all savings') - parser.add_argument('--cuda_device', type=str, default='cuda:0', help='cuda device number, eg:[cuda:0]') - parser.add_argument('--num_layers', type=int, default=2, help='number of layers') - parser.add_argument('--learning_rate', type=float, default=2e-4, help='learning rate') - parser.add_argument('--alpha', type=float, default=0, help='penalty of sparsity') - parser.add_argument('--batch_size', type=int, default=256, help='batch_size') - parser.add_argument('--upsampling', type=str, default='linear', help="'stride2' or 'nearest', 'linear'") - parser.add_argument('--downsampling', type=str, default='stride2', help='stride2 or max_pooling') - parser.add_argument('--batch_normalization', type=int, default=0, help='batch_norm: 1 or 0') - parser.add_argument('--activation', type=str, default='LeakyReLU', help='activation: ReLU, LeakyReLU, tanh') - parser.add_argument('--rotation', type=str, default='quaternion', help='representation of rotation:euler_angle, quaternion') - parser.add_argument('--data_augment', type=int, default=1, help='data_augment: 1 or 0') - parser.add_argument('--epoch_num', type=int, default=20001, help='epoch_num') - parser.add_argument('--window_size', type=int, default=64, help='length of time axis per window') - parser.add_argument('--kernel_size', type=int, default=15, help='must be odd') - parser.add_argument('--base_channel_num', type=int, default=-1) - parser.add_argument('--normalization', type=int, default=1) - parser.add_argument('--verbose', type=int, default=0) - parser.add_argument('--skeleton_dist', type=int, default=2) - parser.add_argument('--skeleton_pool', type=str, default='mean') - parser.add_argument('--extra_conv', type=int, default=0) - parser.add_argument('--padding_mode', type=str, default='reflection') - parser.add_argument('--dataset', type=str, default='Mixamo') - parser.add_argument('--fk_world', type=int, default=0) - parser.add_argument('--patch_gan', type=int, default=1) - parser.add_argument('--debug', type=int, default=0) - parser.add_argument('--skeleton_info', type=str, default='concat') - parser.add_argument('--ee_loss_fact', type=str, default='height') - parser.add_argument('--pos_repr', type=str, default='3d') - parser.add_argument('--D_global_velo', type=int, default=0) - parser.add_argument('--gan_mode', type=str, default='lsgan') - parser.add_argument('--pool_size', type=int, default=50) - parser.add_argument('--is_train', type=int, default=1) - parser.add_argument('--model', type=str, default='mul_top_mul_ske') - parser.add_argument('--epoch_begin', type=int, default=0) - parser.add_argument('--lambda_rec', type=float, default=5) - parser.add_argument('--lambda_cycle', type=float, default=5) - parser.add_argument('--lambda_ee', type=float, default=100) - parser.add_argument('--lambda_global_pose', type=float, default=2.5) - parser.add_argument('--lambda_position', type=float, default=1) - parser.add_argument('--ee_velo', type=int, default=1) - parser.add_argument('--ee_from_root', type=int, default=1) - parser.add_argument('--scheduler', type=str, default='none') - parser.add_argument('--rec_loss_mode', type=str, default='extra_global_pos') - parser.add_argument('--adaptive_ee', type=int, default=0) - parser.add_argument('--simple_operator', type=int, default=0) - parser.add_argument('--use_sep_ee', type=int, default=0) - parser.add_argument('--eval_seq', type=int, default=0) + parser.add_argument( + "--save_dir", type=str, default="./pretrained", help="directory for all savings" + ) + parser.add_argument( + "--cuda_device", + type=str, + default="cuda:0", + help="cuda device number, eg:[cuda:0]", + ) + parser.add_argument("--num_layers", type=int, default=2, help="number of layers") + parser.add_argument( + "--learning_rate", type=float, default=2e-4, help="learning rate" + ) + parser.add_argument("--alpha", type=float, default=0, help="penalty of sparsity") + parser.add_argument("--batch_size", type=int, default=256, help="batch_size") + parser.add_argument( + "--upsampling", + type=str, + default="linear", + help="'stride2' or 'nearest', 'linear'", + ) + parser.add_argument( + "--downsampling", type=str, default="stride2", help="stride2 or max_pooling" + ) + parser.add_argument( + "--batch_normalization", type=int, default=0, help="batch_norm: 1 or 0" + ) + parser.add_argument( + "--activation", + type=str, + default="LeakyReLU", + help="activation: ReLU, LeakyReLU, tanh", + ) + parser.add_argument( + "--rotation", + type=str, + default="quaternion", + help="representation of rotation:euler_angle, quaternion", + ) + parser.add_argument( + "--data_augment", type=int, default=1, help="data_augment: 1 or 0" + ) + parser.add_argument("--epoch_num", type=int, default=20001, help="epoch_num") + parser.add_argument( + "--window_size", type=int, default=64, help="length of time axis per window" + ) + parser.add_argument("--kernel_size", type=int, default=15, help="must be odd") + parser.add_argument("--base_channel_num", type=int, default=-1) + parser.add_argument("--normalization", type=int, default=1) + parser.add_argument("--verbose", type=int, default=0) + parser.add_argument("--skeleton_dist", type=int, default=2) + parser.add_argument("--skeleton_pool", type=str, default="mean") + parser.add_argument("--extra_conv", type=int, default=0) + parser.add_argument("--padding_mode", type=str, default="reflection") + parser.add_argument("--dataset", type=str, default="Mixamo") + parser.add_argument("--fk_world", type=int, default=0) + parser.add_argument("--patch_gan", type=int, default=1) + parser.add_argument("--debug", type=int, default=0) + parser.add_argument("--skeleton_info", type=str, default="concat") + parser.add_argument("--ee_loss_fact", type=str, default="height") + parser.add_argument("--pos_repr", type=str, default="3d") + parser.add_argument("--D_global_velo", type=int, default=0) + parser.add_argument("--gan_mode", type=str, default="lsgan") + parser.add_argument("--pool_size", type=int, default=50) + parser.add_argument("--is_train", type=int, default=1) + parser.add_argument("--model", type=str, default="mul_top_mul_ske") + parser.add_argument("--epoch_begin", type=int, default=0) + parser.add_argument("--lambda_rec", type=float, default=5) + parser.add_argument("--lambda_cycle", type=float, default=5) + parser.add_argument("--lambda_ee", type=float, default=100) + parser.add_argument("--lambda_global_pose", type=float, default=2.5) + parser.add_argument("--lambda_position", type=float, default=1) + parser.add_argument("--ee_velo", type=int, default=1) + parser.add_argument("--ee_from_root", type=int, default=1) + parser.add_argument("--scheduler", type=str, default="none") + parser.add_argument("--rec_loss_mode", type=str, default="extra_global_pos") + parser.add_argument("--adaptive_ee", type=int, default=0) + parser.add_argument("--simple_operator", type=int, default=0) + parser.add_argument("--use_sep_ee", type=int, default=0) + parser.add_argument("--eval_seq", type=int, default=0) return parser @@ -60,13 +92,16 @@ def get_args(): def get_std_bvh(args=None, dataset=None): - if args is None and dataset is None: raise Exception('Unexpected parameter') - if dataset is None: dataset = args.dataset - std_bvh = './datasets/Mixamo/std_bvhs/{}.bvh'.format(dataset) + if args is None and dataset is None: + raise Exception("Unexpected parameter") + if dataset is None: + dataset = args.dataset + std_bvh = "./datasets/Mixamo/std_bvhs/{}.bvh".format(dataset) return std_bvh def try_mkdir(path): import os + if not os.path.exists(path): os.makedirs(path) diff --git a/options/options.py b/options/options.py index 1cc0878..0ab6a92 100644 --- a/options/options.py +++ b/options/options.py @@ -3,24 +3,37 @@ import models -class Options(): - +class Options: def __init__(self): self.initialized = False def initialize(self, parser): - parser.add_argument('--input_A', required=True, help='path to first bvh file') - parser.add_argument('--input_B', required=True, help='path to second bvh file') - parser.add_argument('--model_path', required=True, help='path to second bvh file') - parser.add_argument('--edit_type', type=str, default='retargeting', help='name of the motion editing operation, retargeting or style_transfer') - parser.add_argument('--gpu_ids', type=str, default='0', help='gpu ids: e.g. 0 0,1,2, 0,2. use -1 for CPU') + parser.add_argument("--input_A", required=True, help="path to first bvh file") + parser.add_argument("--input_B", required=True, help="path to second bvh file") + parser.add_argument( + "--model_path", required=True, help="path to second bvh file" + ) + parser.add_argument( + "--edit_type", + type=str, + default="retargeting", + help="name of the motion editing operation, retargeting or style_transfer", + ) + parser.add_argument( + "--gpu_ids", + type=str, + default="0", + help="gpu ids: e.g. 0 0,1,2, 0,2. use -1 for CPU", + ) return parser def gather_options(self): if not self.initialized: # check if it has been initialized - parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) + parser = argparse.ArgumentParser( + formatter_class=argparse.ArgumentDefaultsHelpFormatter + ) parser = self.initialize(parser) # get the basic options @@ -42,30 +55,30 @@ def gather_options(self): return parser.parse_args() def print_options(self, opt): - message = '' - message += '----------------- Options ---------------\n' + message = "" + message += "----------------- Options ---------------\n" for k, v in sorted(vars(opt).items()): - comment = '' + comment = "" default = self.parser.get_default(k) if v != default: - comment = '\t[default: %s]' % str(default) - message += '{:>25}: {:<30}{}\n'.format(str(k), str(v), comment) - message += '----------------- End -------------------' + comment = "\t[default: %s]" % str(default) + message += "{:>25}: {:<30}{}\n".format(str(k), str(v), comment) + message += "----------------- End -------------------" print(message) def parse(self): opt = self.gather_options() - opt.isTrain = self.isTrain # train or test + opt.isTrain = self.isTrain # train or test # process opt.suffix if opt.suffix: - suffix = ('_' + opt.suffix.format(**vars(opt))) if opt.suffix != '' else '' + suffix = ("_" + opt.suffix.format(**vars(opt))) if opt.suffix != "" else "" opt.name = opt.name + suffix self.print_options(opt) # set gpu ids - str_ids = opt.gpu_ids.split(',') + str_ids = opt.gpu_ids.split(",") opt.gpu_ids = [] for str_id in str_ids: id = int(str_id) diff --git a/utils/Animation.py b/utils/Animation.py index 5409b80..1778b51 100644 --- a/utils/Animation.py +++ b/utils/Animation.py @@ -6,6 +6,7 @@ import AnimationStructure from Quaternions_old import Quaternions + class Animation: """ Animation is a numpy-like wrapper for animation data @@ -27,54 +28,74 @@ class Animation: parents : (J) ndarray | Joint Parents """ - + def __init__(self, rotations, positions, orients, offsets, parents): - + self.rotations = rotations self.positions = positions - self.orients = orients - self.offsets = offsets - self.parents = parents - + self.orients = orients + self.offsets = offsets + self.parents = parents + def __op__(self, op, other): return Animation( op(self.rotations, other.rotations), op(self.positions, other.positions), op(self.orients, other.orients), op(self.offsets, other.offsets), - op(self.parents, other.parents)) + op(self.parents, other.parents), + ) def __iop__(self, op, other): self.rotations = op(self.roations, other.rotations) self.positions = op(self.roations, other.positions) - self.orients = op(self.orients, other.orients) - self.offsets = op(self.offsets, other.offsets) - self.parents = op(self.parents, other.parents) + self.orients = op(self.orients, other.orients) + self.offsets = op(self.offsets, other.offsets) + self.parents = op(self.parents, other.parents) return self - + def __sop__(self, op): return Animation( op(self.rotations), op(self.positions), op(self.orients), op(self.offsets), - op(self.parents)) - - def __add__(self, other): return self.__op__(operator.add, other) - def __sub__(self, other): return self.__op__(operator.sub, other) - def __mul__(self, other): return self.__op__(operator.mul, other) - def __div__(self, other): return self.__op__(operator.div, other) - - def __abs__(self): return self.__sop__(operator.abs) - def __neg__(self): return self.__sop__(operator.neg) - - def __iadd__(self, other): return self.__iop__(operator.iadd, other) - def __isub__(self, other): return self.__iop__(operator.isub, other) - def __imul__(self, other): return self.__iop__(operator.imul, other) - def __idiv__(self, other): return self.__iop__(operator.idiv, other) - - def __len__(self): return len(self.rotations) - + op(self.parents), + ) + + def __add__(self, other): + return self.__op__(operator.add, other) + + def __sub__(self, other): + return self.__op__(operator.sub, other) + + def __mul__(self, other): + return self.__op__(operator.mul, other) + + def __div__(self, other): + return self.__op__(operator.div, other) + + def __abs__(self): + return self.__sop__(operator.abs) + + def __neg__(self): + return self.__sop__(operator.neg) + + def __iadd__(self, other): + return self.__iop__(operator.iadd, other) + + def __isub__(self, other): + return self.__iop__(operator.isub, other) + + def __imul__(self, other): + return self.__iop__(operator.imul, other) + + def __idiv__(self, other): + return self.__iop__(operator.idiv, other) + + def __len__(self): + return len(self.rotations) + def __getitem__(self, k): if isinstance(k, tuple): return Animation( @@ -82,16 +103,18 @@ def __getitem__(self, k): self.positions[k], self.orients[k[1:]], self.offsets[k[1:]], - self.parents[k[1:]]) + self.parents[k[1:]], + ) else: return Animation( self.rotations[k], self.positions[k], self.orients, self.offsets, - self.parents) - - def __setitem__(self, k, v): + self.parents, + ) + + def __setitem__(self, k, v): if isinstance(k, tuple): self.rotations.__setitem__(k, v.rotations) self.positions.__setitem__(k, v.positions) @@ -104,43 +127,58 @@ def __setitem__(self, k, v): self.orients.__setitem__(k, v.orients) self.offsets.__setitem__(k, v.offsets) self.parents.__setitem__(k, v.parents) - + @property - def shape(self): return (self.rotations.shape[0], self.rotations.shape[1]) - - def copy(self): return Animation( - self.rotations.copy(), self.positions.copy(), - self.orients.copy(), self.offsets.copy(), - self.parents.copy()) - + def shape(self): + return (self.rotations.shape[0], self.rotations.shape[1]) + + def copy(self): + return Animation( + self.rotations.copy(), + self.positions.copy(), + self.orients.copy(), + self.offsets.copy(), + self.parents.copy(), + ) + def repeat(self, *args, **kw): return Animation( self.rotations.repeat(*args, **kw), self.positions.repeat(*args, **kw), - self.orients, self.offsets, self.parents) - + self.orients, + self.offsets, + self.parents, + ) + def ravel(self): - return np.hstack([ - self.rotations.log().ravel(), - self.positions.ravel(), - self.orients.log().ravel(), - self.offsets.ravel()]) - + return np.hstack( + [ + self.rotations.log().ravel(), + self.positions.ravel(), + self.orients.log().ravel(), + self.offsets.ravel(), + ] + ) + @classmethod def unravel(clas, anim, shape, parents): nf, nj = shape - rotations = anim[nf*nj*0:nf*nj*3] - positions = anim[nf*nj*3:nf*nj*6] - orients = anim[nf*nj*6+nj*0:nf*nj*6+nj*3] - offsets = anim[nf*nj*6+nj*3:nf*nj*6+nj*6] + rotations = anim[nf * nj * 0 : nf * nj * 3] + positions = anim[nf * nj * 3 : nf * nj * 6] + orients = anim[nf * nj * 6 + nj * 0 : nf * nj * 6 + nj * 3] + offsets = anim[nf * nj * 6 + nj * 3 : nf * nj * 6 + nj * 6] return cls( - Quaternions.exp(rotations), positions, - Quaternions.exp(orients), offsets, - parents.copy()) - - + Quaternions.exp(rotations), + positions, + Quaternions.exp(orients), + offsets, + parents.copy(), + ) + + """ Maya Interaction """ + def load_to_maya(anim, names=None, radius=0.5): """ Load Animation Object into Maya as Joint Skeleton @@ -165,53 +203,57 @@ def load_to_maya(anim, names=None, radius=0.5): List of Maya Joint Nodes loaded into scene """ - + import pymel.core as pm - + joints = [] - frames = range(1, len(anim)+1) - - if names is None: names = ["joint_" + str(i) for i in range(len(anim.parents))] - - for i, offset, orient, parent, name in zip(range(len(anim.offsets)), anim.offsets, anim.orients, anim.parents, names): - + frames = range(1, len(anim) + 1) + + if names is None: + names = ["joint_" + str(i) for i in range(len(anim.parents))] + + for i, offset, orient, parent, name in zip( + range(len(anim.offsets)), anim.offsets, anim.orients, anim.parents, names + ): + if parent < 0: pm.select(d=True) else: pm.select(joints[parent]) - + joint = pm.joint(n=name, p=offset, relative=True, radius=radius) joint.setOrientation([orient[1], orient[2], orient[3], orient[0]]) - + curvex = pm.nodetypes.AnimCurveTA(n=name + "_rotateX") curvey = pm.nodetypes.AnimCurveTA(n=name + "_rotateY") - curvez = pm.nodetypes.AnimCurveTA(n=name + "_rotateZ") - - jrotations = (-Quaternions(orient[np.newaxis]) * anim.rotations[:,i]).euler() - curvex.addKeys(frames, jrotations[:,0]) - curvey.addKeys(frames, jrotations[:,1]) - curvez.addKeys(frames, jrotations[:,2]) - + curvez = pm.nodetypes.AnimCurveTA(n=name + "_rotateZ") + + jrotations = (-Quaternions(orient[np.newaxis]) * anim.rotations[:, i]).euler() + curvex.addKeys(frames, jrotations[:, 0]) + curvey.addKeys(frames, jrotations[:, 1]) + curvez.addKeys(frames, jrotations[:, 2]) + pm.connectAttr(curvex.output, joint.rotateX) pm.connectAttr(curvey.output, joint.rotateY) pm.connectAttr(curvez.output, joint.rotateZ) - + offsetx = pm.nodetypes.AnimCurveTU(n=name + "_translateX") offsety = pm.nodetypes.AnimCurveTU(n=name + "_translateY") offsetz = pm.nodetypes.AnimCurveTU(n=name + "_translateZ") - - offsetx.addKeys(frames, anim.positions[:,i,0]) - offsety.addKeys(frames, anim.positions[:,i,1]) - offsetz.addKeys(frames, anim.positions[:,i,2]) - + + offsetx.addKeys(frames, anim.positions[:, i, 0]) + offsety.addKeys(frames, anim.positions[:, i, 1]) + offsetz.addKeys(frames, anim.positions[:, i, 2]) + pm.connectAttr(offsetx.output, joint.translateX) pm.connectAttr(offsety.output, joint.translateY) pm.connectAttr(offsetz.output, joint.translateZ) - + joints.append(joint) - + return joints + def load_from_maya(root, start, end): """ Load Animation Object from Maya Joint Skeleton @@ -236,60 +278,68 @@ def load_from_maya(root, start, end): """ import pymel.core as pm - + original_time = pm.currentTime(q=True) pm.currentTime(start) - + """ Build Structure """ - + names, parents = AnimationStructure.load_from_maya(root) descendants = AnimationStructure.descendants_list(parents) orients = Quaternions.id(len(names)) offsets = np.array([pm.xform(j, q=True, translation=True) for j in names]) - + for j, name in enumerate(names): scale = pm.xform(pm.PyNode(name), q=True, scale=True, relative=True) - if len(descendants[j]) == 0: continue + if len(descendants[j]) == 0: + continue offsets[descendants[j]] *= scale - + """ Load Animation """ - eulers = np.zeros((end-start, len(names), 3)) - positions = np.zeros((end-start, len(names), 3)) - rotations = Quaternions.id((end-start, len(names))) - - for i in range(end-start): - - pm.currentTime(start+i+1, u=True) - + eulers = np.zeros((end - start, len(names), 3)) + positions = np.zeros((end - start, len(names), 3)) + rotations = Quaternions.id((end - start, len(names))) + + for i in range(end - start): + + pm.currentTime(start + i + 1, u=True) + scales = {} - + for j, name, parent in zip(range(len(names)), names, parents): - + node = pm.PyNode(name) - - if i == 0 and pm.hasAttr(node, 'jointOrient'): + + if i == 0 and pm.hasAttr(node, "jointOrient"): ort = node.getOrientation() orients[j] = Quaternions(np.array([ort[3], ort[0], ort[1], ort[2]])) - - if pm.hasAttr(node, 'rotate'): eulers[i,j] = np.radians(pm.xform(node, q=True, rotation=True)) - if pm.hasAttr(node, 'translate'): positions[i,j] = pm.xform(node, q=True, translation=True) - if pm.hasAttr(node, 'scale'): scales[j] = pm.xform(node, q=True, scale=True, relative=True) + + if pm.hasAttr(node, "rotate"): + eulers[i, j] = np.radians(pm.xform(node, q=True, rotation=True)) + if pm.hasAttr(node, "translate"): + positions[i, j] = pm.xform(node, q=True, translation=True) + if pm.hasAttr(node, "scale"): + scales[j] = pm.xform(node, q=True, scale=True, relative=True) for j in scales: - if len(descendants[j]) == 0: continue - positions[i,descendants[j]] *= scales[j] - - positions[i,0] = pm.xform(root, q=True, translation=True, worldSpace=True) - - rotations = orients[np.newaxis] * Quaternions.from_euler(eulers, order='xyz', world=True) - + if len(descendants[j]) == 0: + continue + positions[i, descendants[j]] *= scales[j] + + positions[i, 0] = pm.xform(root, q=True, translation=True, worldSpace=True) + + rotations = orients[np.newaxis] * Quaternions.from_euler( + eulers, order="xyz", world=True + ) + """ Done """ - + pm.currentTime(original_time) - + return Animation(rotations, positions, orients, offsets, parents), names - + + # local transformation matrices def transforms_local(anim): """ @@ -314,16 +364,20 @@ def transforms_local(anim): For each frame F, joint local transforms for each joint J """ - + transforms = anim.rotations.transforms() - transforms = np.concatenate([transforms, np.zeros(transforms.shape[:2] + (3, 1))], axis=-1) - transforms = np.concatenate([transforms, np.zeros(transforms.shape[:2] + (1, 4))], axis=-2) + transforms = np.concatenate( + [transforms, np.zeros(transforms.shape[:2] + (3, 1))], axis=-1 + ) + transforms = np.concatenate( + [transforms, np.zeros(transforms.shape[:2] + (1, 4))], axis=-2 + ) # the last column is filled with the joint positions! - transforms[:,:,0:3,3] = anim.positions - transforms[:,:,3:4,3] = 1.0 + transforms[:, :, 0:3, 3] = anim.positions + transforms[:, :, 3:4, 3] = 1.0 return transforms - + def transforms_multiply(t0s, t1s): """ Transforms Multiply @@ -346,14 +400,16 @@ def transforms_multiply(t0s, t1s): frame F and joint J multiplied together """ - + return ut.matrix_multiply(t0s, t1s) - + + def transforms_inv(ts): fts = ts.reshape(-1, 4, 4) fts = np.array(list(map(lambda x: np.linalg.inv(x), fts))) return fts.reshape(ts.shape) - + + def transforms_blank(anim): """ Blank Transforms @@ -372,11 +428,14 @@ def transforms_blank(anim): each frame F and joint J """ - ts = np.zeros(anim.shape + (4, 4)) - ts[:,:,0,0] = 1.0; ts[:,:,1,1] = 1.0; - ts[:,:,2,2] = 1.0; ts[:,:,3,3] = 1.0; + ts = np.zeros(anim.shape + (4, 4)) + ts[:, :, 0, 0] = 1.0 + ts[:, :, 1, 1] = 1.0 + ts[:, :, 2, 2] = 1.0 + ts[:, :, 3, 3] = 1.0 return ts + # global transformation matrices def transforms_global(anim): """ @@ -401,19 +460,20 @@ def transforms_global(anim): Array of global transforms for each frame F and joint J """ - - joints = np.arange(anim.shape[1]) + + joints = np.arange(anim.shape[1]) parents = np.arange(anim.shape[1]) - locals = transforms_local(anim) + locals = transforms_local(anim) globals = transforms_blank(anim) - globals[:,0] = locals[:,0] - + globals[:, 0] = locals[:, 0] + for i in range(1, anim.shape[1]): - globals[:,i] = transforms_multiply(globals[:,anim.parents[i]], locals[:,i]) - + globals[:, i] = transforms_multiply(globals[:, anim.parents[i]], locals[:, i]) + return globals - + + # !!! useful! def positions_global(anim): """ @@ -437,11 +497,13 @@ def positions_global(anim): """ # get the last column -- corresponding to the coordinates - positions = transforms_global(anim)[:,:,:,3] - return positions[:,:,:3] / positions[:,:,3,np.newaxis] - + positions = transforms_global(anim)[:, :, :, 3] + return positions[:, :, :3] / positions[:, :, 3, np.newaxis] + + """ Rotations """ - + + def rotations_global(anim): """ Global Animation Rotations @@ -466,24 +528,26 @@ def rotations_global(anim): and joint J """ - joints = np.arange(anim.shape[1]) + joints = np.arange(anim.shape[1]) parents = np.arange(anim.shape[1]) - locals = anim.rotations + locals = anim.rotations globals = Quaternions.id(anim.shape) - - globals[:,0] = locals[:,0] - + + globals[:, 0] = locals[:, 0] + for i in range(1, anim.shape[1]): - globals[:,i] = globals[:,anim.parents[i]] * locals[:,i] - + globals[:, i] = globals[:, anim.parents[i]] * locals[:, i] + return globals - + + def rotations_parents_global(anim): rotations = rotations_global(anim) - rotations = rotations[:,anim.parents] - rotations[:,0] = Quaternions.id(len(anim)) + rotations = rotations[:, anim.parents] + rotations[:, 0] = Quaternions.id(len(anim)) return rotations - + + def rotations_load_to_maya(rotations, positions, names=None): """ Load Rotations into Maya @@ -515,130 +579,142 @@ def rotations_load_to_maya(rotations, positions, names=None): maxies : Group Grouped Maya Node of all Axis nodes """ - + import pymel.core as pm - if names is None: names = ["joint_" + str(i) for i in range(rotations.shape[1])] - + if names is None: + names = ["joint_" + str(i) for i in range(rotations.shape[1])] + maxis = [] - frames = range(1, len(positions)+1) + frames = range(1, len(positions) + 1) for i, name in enumerate(names): - + name = name + "_axis" axis = pm.group( - pm.curve(p=[(0,0,0), (1,0,0)], d=1, n=name+'_axis_x'), - pm.curve(p=[(0,0,0), (0,1,0)], d=1, n=name+'_axis_y'), - pm.curve(p=[(0,0,0), (0,0,1)], d=1, n=name+'_axis_z'), - n=name) - - axis.rotatePivot.set((0,0,0)) - axis.scalePivot.set((0,0,0)) - axis.childAtIndex(0).overrideEnabled.set(1); axis.childAtIndex(0).overrideColor.set(13) - axis.childAtIndex(1).overrideEnabled.set(1); axis.childAtIndex(1).overrideColor.set(14) - axis.childAtIndex(2).overrideEnabled.set(1); axis.childAtIndex(2).overrideColor.set(15) - + pm.curve(p=[(0, 0, 0), (1, 0, 0)], d=1, n=name + "_axis_x"), + pm.curve(p=[(0, 0, 0), (0, 1, 0)], d=1, n=name + "_axis_y"), + pm.curve(p=[(0, 0, 0), (0, 0, 1)], d=1, n=name + "_axis_z"), + n=name, + ) + + axis.rotatePivot.set((0, 0, 0)) + axis.scalePivot.set((0, 0, 0)) + axis.childAtIndex(0).overrideEnabled.set(1) + axis.childAtIndex(0).overrideColor.set(13) + axis.childAtIndex(1).overrideEnabled.set(1) + axis.childAtIndex(1).overrideColor.set(14) + axis.childAtIndex(2).overrideEnabled.set(1) + axis.childAtIndex(2).overrideColor.set(15) + curvex = pm.nodetypes.AnimCurveTA(n=name + "_rotateX") curvey = pm.nodetypes.AnimCurveTA(n=name + "_rotateY") - curvez = pm.nodetypes.AnimCurveTA(n=name + "_rotateZ") - - arotations = rotations[:,i].euler() - curvex.addKeys(frames, arotations[:,0]) - curvey.addKeys(frames, arotations[:,1]) - curvez.addKeys(frames, arotations[:,2]) - + curvez = pm.nodetypes.AnimCurveTA(n=name + "_rotateZ") + + arotations = rotations[:, i].euler() + curvex.addKeys(frames, arotations[:, 0]) + curvey.addKeys(frames, arotations[:, 1]) + curvez.addKeys(frames, arotations[:, 2]) + pm.connectAttr(curvex.output, axis.rotateX) pm.connectAttr(curvey.output, axis.rotateY) pm.connectAttr(curvez.output, axis.rotateZ) - + offsetx = pm.nodetypes.AnimCurveTU(n=name + "_translateX") offsety = pm.nodetypes.AnimCurveTU(n=name + "_translateY") offsetz = pm.nodetypes.AnimCurveTU(n=name + "_translateZ") - - offsetx.addKeys(frames, positions[:,i,0]) - offsety.addKeys(frames, positions[:,i,1]) - offsetz.addKeys(frames, positions[:,i,2]) - + + offsetx.addKeys(frames, positions[:, i, 0]) + offsety.addKeys(frames, positions[:, i, 1]) + offsetz.addKeys(frames, positions[:, i, 2]) + pm.connectAttr(offsetx.output, axis.translateX) pm.connectAttr(offsety.output, axis.translateY) pm.connectAttr(offsetz.output, axis.translateZ) - + maxis.append(axis) - - return pm.group(*maxis, n='RotationAnimation') - + + return pm.group(*maxis, n="RotationAnimation") + + """ Offsets & Orients """ + def orients_global(anim): - joints = np.arange(anim.shape[1]) + joints = np.arange(anim.shape[1]) parents = np.arange(anim.shape[1]) - locals = anim.orients + locals = anim.orients globals = Quaternions.id(anim.shape[1]) - - globals[:,0] = locals[:,0] - + + globals[:, 0] = locals[:, 0] + for i in range(1, anim.shape[1]): - globals[:,i] = globals[:,anim.parents[i]] * locals[:,i] - + globals[:, i] = globals[:, anim.parents[i]] * locals[:, i] + return globals - + def offsets_transforms_local(anim): - + transforms = anim.orients[np.newaxis].transforms() - transforms = np.concatenate([transforms, np.zeros(transforms.shape[:2] + (3, 1))], axis=-1) - transforms = np.concatenate([transforms, np.zeros(transforms.shape[:2] + (1, 4))], axis=-2) - transforms[:,:,0:3,3] = anim.offsets[np.newaxis] - transforms[:,:,3:4,3] = 1.0 + transforms = np.concatenate( + [transforms, np.zeros(transforms.shape[:2] + (3, 1))], axis=-1 + ) + transforms = np.concatenate( + [transforms, np.zeros(transforms.shape[:2] + (1, 4))], axis=-2 + ) + transforms[:, :, 0:3, 3] = anim.offsets[np.newaxis] + transforms[:, :, 3:4, 3] = 1.0 return transforms - - + + def offsets_transforms_global(anim): - - joints = np.arange(anim.shape[1]) + + joints = np.arange(anim.shape[1]) parents = np.arange(anim.shape[1]) - locals = offsets_transforms_local(anim) + locals = offsets_transforms_local(anim) globals = transforms_blank(anim) - globals[:,0] = locals[:,0] - + globals[:, 0] = locals[:, 0] + for i in range(1, anim.shape[1]): - globals[:,i] = transforms_multiply(globals[:,anim.parents[i]], locals[:,i]) - + globals[:, i] = transforms_multiply(globals[:, anim.parents[i]], locals[:, i]) + return globals - + + def offsets_global(anim): - offsets = offsets_transforms_global(anim)[:,:,:,3] - return offsets[0,:,:3] / offsets[0,:,3,np.newaxis] - + offsets = offsets_transforms_global(anim)[:, :, :, 3] + return offsets[0, :, :3] / offsets[0, :, 3, np.newaxis] + + """ Lengths """ + def offset_lengths(anim): - return np.sum(anim.offsets[1:]**2.0, axis=1)**0.5 - - + return np.sum(anim.offsets[1:] ** 2.0, axis=1) ** 0.5 + + def position_lengths(anim): - return np.sum(anim.positions[:,1:]**2.0, axis=2)**0.5 - - + return np.sum(anim.positions[:, 1:] ** 2.0, axis=2) ** 0.5 + + """ Skinning """ + def skin(anim, rest, weights, mesh, maxjoints=4): - + full_transforms = transforms_multiply( - transforms_global(anim), - transforms_inv(transforms_global(rest[0:1]))) - - weightids = np.argsort(-weights, axis=1)[:,:maxjoints] - weightvls = np.array(list(map(lambda w, i: w[i], weights, weightids))) - weightvls = weightvls / weightvls.sum(axis=1)[...,np.newaxis] - - verts = np.hstack([mesh, np.ones((len(mesh), 1))]) - verts = verts[np.newaxis,:,np.newaxis,:,np.newaxis] - verts = transforms_multiply(full_transforms[:,weightids], verts) - verts = (verts[:,:,:,:3] / verts[:,:,:,3:4])[:,:,:,:,0] + transforms_global(anim), transforms_inv(transforms_global(rest[0:1])) + ) - return np.sum(weightvls[np.newaxis,:,:,np.newaxis] * verts, axis=2) - + weightids = np.argsort(-weights, axis=1)[:, :maxjoints] + weightvls = np.array(list(map(lambda w, i: w[i], weights, weightids))) + weightvls = weightvls / weightvls.sum(axis=1)[..., np.newaxis] + verts = np.hstack([mesh, np.ones((len(mesh), 1))]) + verts = verts[np.newaxis, :, np.newaxis, :, np.newaxis] + verts = transforms_multiply(full_transforms[:, weightids], verts) + verts = (verts[:, :, :, :3] / verts[:, :, :, 3:4])[:, :, :, :, 0] + return np.sum(weightvls[np.newaxis, :, :, np.newaxis] * verts, axis=2) diff --git a/utils/AnimationStructure.py b/utils/AnimationStructure.py index 28439df..5d83195 100644 --- a/utils/AnimationStructure.py +++ b/utils/AnimationStructure.py @@ -5,6 +5,7 @@ """ Maya Functions """ + def load_from_maya(root): """ Load joint parents and names from maya @@ -26,33 +27,40 @@ def load_from_maya(root): Joint index -1 is used to represent that there is no parent joint """ - + import pymel.core as pm - + names = [] parents = [] def unload_joint(j, parents, par): - + id = len(names) names.append(j) parents.append(par) - - children = [c for c in j.getChildren() if - isinstance(c, pm.nt.Transform) and - not isinstance(c, pm.nt.Constraint) and - not any(pm.listRelatives(c, s=True)) and - (any(pm.listRelatives(c, ad=True, ap=False, type='joint')) or isinstance(c, pm.nt.Joint))] - + + children = [ + c + for c in j.getChildren() + if isinstance(c, pm.nt.Transform) + and not isinstance(c, pm.nt.Constraint) + and not any(pm.listRelatives(c, s=True)) + and ( + any(pm.listRelatives(c, ad=True, ap=False, type="joint")) + or isinstance(c, pm.nt.Joint) + ) + ] + map(lambda c: unload_joint(c, parents, id), children) unload_joint(root, parents, -1) - + return (names, parents) - + """ Family Functions """ + def joints(parents): """ Parameters @@ -69,6 +77,7 @@ def joints(parents): """ return np.arange(len(parents), dtype=int) + def joints_list(parents): """ Parameters @@ -84,8 +93,9 @@ def joints_list(parents): List of arrays of joint idices for each joint """ - return list(joints(parents)[:,np.newaxis]) - + return list(joints(parents)[:, np.newaxis]) + + def parents_list(parents): """ Parameters @@ -101,9 +111,9 @@ def parents_list(parents): List of arrays of joint idices for the parents of each joint """ - return list(parents[:,np.newaxis]) - - + return list(parents[:, np.newaxis]) + + def children_list(parents): """ Parameters @@ -119,13 +129,13 @@ def children_list(parents): List of arrays of joint indices for the children of each joint """ - + def joint_children(i): return [j for j, p in enumerate(parents) if p == i] - + return list(map(lambda j: np.array(joint_children(j)), joints(parents))) - - + + def descendants_list(parents): """ Parameters @@ -141,15 +151,15 @@ def descendants_list(parents): List of arrays of joint idices for the descendants of each joint """ - + children = children_list(parents) - + def joint_descendants(i): return sum([joint_descendants(j) for j in children[i]], list(children[i])) - + return list(map(lambda j: np.array(joint_descendants(j)), joints(parents))) - - + + def ancestors_list(parents): """ Parameters @@ -165,17 +175,18 @@ def ancestors_list(parents): List of arrays of joint idices for the ancestors of each joint """ - + decendants = descendants_list(parents) - + def joint_ancestors(i): return [j for j in joints(parents) if i in decendants[j]] - + return list(map(lambda j: np.array(joint_ancestors(j)), joints(parents))) - - + + """ Mask Functions """ + def mask(parents, filter): """ Constructs a Mask for a give filter @@ -208,17 +219,34 @@ def mask(parents, filter): m = np.zeros((len(parents), len(parents))).astype(bool) jnts = joints(parents) fltr = filter(parents) - for i,f in enumerate(fltr): m[i,:] = np.any(jnts[:,np.newaxis] == f[np.newaxis,:], axis=1) + for i, f in enumerate(fltr): + m[i, :] = np.any(jnts[:, np.newaxis] == f[np.newaxis, :], axis=1) return m -def joints_mask(parents): return np.eye(len(parents)).astype(bool) -def children_mask(parents): return mask(parents, children_list) -def parents_mask(parents): return mask(parents, parents_list) -def descendants_mask(parents): return mask(parents, descendants_list) -def ancestors_mask(parents): return mask(parents, ancestors_list) - + +def joints_mask(parents): + return np.eye(len(parents)).astype(bool) + + +def children_mask(parents): + return mask(parents, children_list) + + +def parents_mask(parents): + return mask(parents, parents_list) + + +def descendants_mask(parents): + return mask(parents, descendants_list) + + +def ancestors_mask(parents): + return mask(parents, ancestors_list) + + """ Search Functions """ - + + def joint_chain_ascend(parents, start, end): chain = [] while start != end: @@ -226,10 +254,11 @@ def joint_chain_ascend(parents, start, end): start = parents[start] chain.append(end) return np.array(chain, dtype=int) - - + + """ Constraints """ - + + def constraints(anim, **kwargs): """ Constraint list for Animation @@ -258,38 +287,41 @@ def constraints(anim, **kwargs): (Joint1, Joint2, Masses1, Masses2, Lengths) """ - - masses = kwargs.pop('masses', None) - + + masses = kwargs.pop("masses", None) + children = children_list(anim.parents) constraints = [] points_offsets = Animation.offsets_global(anim) points = Animation.positions_global(anim) - + if masses is None: - masses = 1.0 / (0.1 + np.absolute(points_offsets[:,1])) + masses = 1.0 / (0.1 + np.absolute(points_offsets[:, 1])) masses = masses[np.newaxis].repeat(len(anim), axis=0) - + for j in xrange(anim.shape[1]): - + """ Add constraints between all joints and their children """ for c0 in children[j]: - - dists = np.sum((points[:, c0] - points[:, j])**2.0, axis=1)**0.5 - constraints.append((c0, j, masses[:,c0], masses[:,j], dists)) - + + dists = np.sum((points[:, c0] - points[:, j]) ** 2.0, axis=1) ** 0.5 + constraints.append((c0, j, masses[:, c0], masses[:, j], dists)) + """ Add constraints between all children of joint """ for c1 in children[j]: - if c0 == c1: continue + if c0 == c1: + continue + + dists = np.sum((points[:, c0] - points[:, c1]) ** 2.0, axis=1) ** 0.5 + constraints.append((c0, c1, masses[:, c0], masses[:, c1], dists)) - dists = np.sum((points[:, c0] - points[:, c1])**2.0, axis=1)**0.5 - constraints.append((c0, c1, masses[:,c0], masses[:,c1], dists)) - return constraints - + + """ Graph Functions """ + def graph(anim): """ Generates a weighted adjacency matrix @@ -324,18 +356,19 @@ def graph(anim): directly connected are assigned the weight `0`. """ - + graph = np.zeros(anim.shape[1], anim.shape[1]) - lengths = np.sum(anim.offsets**2.0, axis=1)**0.5 + 0.001 - - for i,p in enumerate(anim.parents): - if p == -1: continue - graph[i,p] = lengths[p] - graph[p,i] = lengths[p] - + lengths = np.sum(anim.offsets ** 2.0, axis=1) ** 0.5 + 0.001 + + for i, p in enumerate(anim.parents): + if p == -1: + continue + graph[i, p] = lengths[p] + graph[p, i] = lengths[p] + return graph - + def distances(anim): """ Generates a distance matrix for @@ -357,43 +390,58 @@ def distances(anim): from some joint N to some joint M """ - + distances = np.zeros((anim.shape[1], anim.shape[1])) generated = distances.copy().astype(bool) - - joint_lengths = np.sum(anim.offsets**2.0, axis=1)**0.5 + + joint_lengths = np.sum(anim.offsets ** 2.0, axis=1) ** 0.5 joint_children = children_list(anim) - joint_parents = parents_list(anim) - + joint_parents = parents_list(anim) + def find_distance(distances, generated, prev, i, j): - - """ If root, identity, or already generated, return """ - if j == -1: return (0.0, True) - if j == i: return (0.0, True) - if generated[i,j]: return (distances[i,j], True) - + + """ If root, identity, or already generated, return """ + if j == -1: + return (0.0, True) + if j == i: + return (0.0, True) + if generated[i, j]: + return (distances[i, j], True) + """ Find best distances along parents and children """ - par_dists = [(joint_lengths[j], find_distance(distances, generated, j, i, p)) for p in joint_parents[j] if p != prev] - out_dists = [(joint_lengths[c], find_distance(distances, generated, j, i, c)) for c in joint_children[j] if c != prev] - + par_dists = [ + (joint_lengths[j], find_distance(distances, generated, j, i, p)) + for p in joint_parents[j] + if p != prev + ] + out_dists = [ + (joint_lengths[c], find_distance(distances, generated, j, i, c)) + for c in joint_children[j] + if c != prev + ] + """ Check valid distance and not dead end """ par_dists = [a + d for (a, (d, f)) in par_dists if f] out_dists = [a + d for (a, (d, f)) in out_dists if f] - + """ All dead ends """ - if (out_dists + par_dists) == []: return (0.0, False) - + if (out_dists + par_dists) == []: + return (0.0, False) + """ Get minimum path """ dist = min(out_dists + par_dists) - distances[i,j] = dist; distances[j,i] = dist - generated[i,j] = True; generated[j,i] = True - + distances[i, j] = dist + distances[j, i] = dist + generated[i, j] = True + generated[j, i] = True + for i in xrange(anim.shape[1]): for j in xrange(anim.shape[1]): find_distance(distances, generated, -1, i, j) - + return distances - + + def edges(parents): """ Animation structure edges @@ -413,10 +461,10 @@ def edges(parents): which corrisponds to an edge in the joint structure going from parent to child. """ - + return np.array(list(zip(parents, joints(parents)))[1:]) - - + + def incidence(parents): """ Incidence Matrix @@ -439,13 +487,12 @@ def incidence(parents): array of vectors along each edge of the structure """ - + es = edges(parents) - - inc = np.zeros((len(parents)-1, len(parents))).astype(np.int) + + inc = np.zeros((len(parents) - 1, len(parents))).astype(np.int) for i, e in enumerate(es): - inc[i,e[0]] = 1 - inc[i,e[1]] = -1 - + inc[i, e[0]] = 1 + inc[i, e[1]] = -1 + return inc.T - \ No newline at end of file diff --git a/utils/BVH.py b/utils/BVH.py index 0801bda..91c691e 100644 --- a/utils/BVH.py +++ b/utils/BVH.py @@ -1,29 +1,27 @@ import re import numpy as np import sys + sys.path.append("motion_utils") from Animation import Animation from Quaternions_old import Quaternions -channelmap = { - 'Xrotation' : 'x', - 'Yrotation' : 'y', - 'Zrotation' : 'z' -} +channelmap = {"Xrotation": "x", "Yrotation": "y", "Zrotation": "z"} channelmap_inv = { - 'x': 'Xrotation', - 'y': 'Yrotation', - 'z': 'Zrotation', + "x": "Xrotation", + "y": "Yrotation", + "z": "Zrotation", } ordermap = { - 'x' : 0, - 'y' : 1, - 'z' : 2, + "x": 0, + "y": 1, + "z": 2, } + def load(filename, start=None, end=None, order=None, world=False): """ Reads a BVH file and constructs an animation @@ -54,52 +52,59 @@ def load(filename, start=None, end=None, order=None, world=False): (animation, joint_names, frametime) Tuple of loaded animation and joint names """ - + f = open(filename, "r") i = 0 active = -1 end_site = False - + names = [] orients = Quaternions.id(0) - offsets = np.array([]).reshape((0,3)) + offsets = np.array([]).reshape((0, 3)) parents = np.array([], dtype=int) - + for line in f: - - if "HIERARCHY" in line: continue - if "MOTION" in line: continue + + if "HIERARCHY" in line: + continue + if "MOTION" in line: + continue rmatch = re.match(r"ROOT (\w+)", line) if rmatch: names.append(rmatch.group(1)) - offsets = np.append(offsets, np.array([[0,0,0]]), axis=0) - orients.qs = np.append(orients.qs, np.array([[1,0,0,0]]), axis=0) - parents = np.append(parents, active) - active = (len(parents)-1) + offsets = np.append(offsets, np.array([[0, 0, 0]]), axis=0) + orients.qs = np.append(orients.qs, np.array([[1, 0, 0, 0]]), axis=0) + parents = np.append(parents, active) + active = len(parents) - 1 continue - if "{" in line: continue + if "{" in line: + continue if "}" in line: - if end_site: end_site = False - else: active = parents[active] + if end_site: + end_site = False + else: + active = parents[active] continue - - offmatch = re.match(r"\s*OFFSET\s+([\-\d\.e]+)\s+([\-\d\.e]+)\s+([\-\d\.e]+)", line) + + offmatch = re.match( + r"\s*OFFSET\s+([\-\d\.e]+)\s+([\-\d\.e]+)\s+([\-\d\.e]+)", line + ) if offmatch: if not end_site: offsets[active] = np.array([list(map(float, offmatch.groups()))]) continue - + chanmatch = re.match(r"\s*CHANNELS\s+(\d+)", line) if chanmatch: channels = int(chanmatch.group(1)) if order is None: channelis = 0 if channels == 3 else 3 channelie = 3 if channels == 3 else 6 - parts = line.split()[2+channelis:2+channelie] + parts = line.split()[2 + channelis : 2 + channelie] if any([p not in channelmap for p in parts]): continue order = "".join([channelmap[p] for p in parts]) @@ -108,20 +113,20 @@ def load(filename, start=None, end=None, order=None, world=False): jmatch = re.match("\s*JOINT\s+(\w+)", line) if jmatch: names.append(jmatch.group(1)) - offsets = np.append(offsets, np.array([[0,0,0]]), axis=0) - orients.qs = np.append(orients.qs, np.array([[1,0,0,0]]), axis=0) - parents = np.append(parents, active) - active = (len(parents)-1) + offsets = np.append(offsets, np.array([[0, 0, 0]]), axis=0) + orients.qs = np.append(orients.qs, np.array([[1, 0, 0, 0]]), axis=0) + parents = np.append(parents, active) + active = len(parents) - 1 continue - + if "End Site" in line: end_site = True continue - + fmatch = re.match("\s*Frames:\s+(\d+)", line) if fmatch: if start and end: - fnum = (end - start)-1 + fnum = (end - start) - 1 else: fnum = int(fmatch.group(1)) jnum = len(parents) @@ -130,46 +135,51 @@ def load(filename, start=None, end=None, order=None, world=False): # result: [fnum, len(orients), 3] rotations = np.zeros((fnum, len(orients), 3)) continue - + fmatch = re.match("\s*Frame Time:\s+([\d\.]+)", line) if fmatch: frametime = float(fmatch.group(1)) continue - - if (start and end) and (i < start or i >= end-1): + + if (start and end) and (i < start or i >= end - 1): i += 1 continue - + dmatch = line.strip().split() if dmatch: data_block = np.array(list(map(float, dmatch))) N = len(parents) fi = i - start if start else i - if channels == 3: + if channels == 3: # This should be root positions[0:1] & all rotations - positions[fi,0:1] = data_block[0:3] - rotations[fi, : ] = data_block[3: ].reshape(N,3) + positions[fi, 0:1] = data_block[0:3] + rotations[fi, :] = data_block[3:].reshape(N, 3) elif channels == 6: - data_block = data_block.reshape(N,6) + data_block = data_block.reshape(N, 6) # fill in all positions - positions[fi,:] = data_block[:,0:3] - rotations[fi,:] = data_block[:,3:6] + positions[fi, :] = data_block[:, 0:3] + rotations[fi, :] = data_block[:, 3:6] elif channels == 9: - positions[fi,0] = data_block[0:3] - data_block = data_block[3:].reshape(N-1,9) - rotations[fi,1:] = data_block[:,3:6] - positions[fi,1:] += data_block[:,0:3] * data_block[:,6:9] + positions[fi, 0] = data_block[0:3] + data_block = data_block[3:].reshape(N - 1, 9) + rotations[fi, 1:] = data_block[:, 3:6] + positions[fi, 1:] += data_block[:, 0:3] * data_block[:, 6:9] else: raise Exception("Too many channels! %i" % channels) i += 1 f.close() - + rotations = Quaternions.from_euler(np.radians(rotations), order=order, world=world) - - return (Animation(rotations, positions, orients, offsets, parents), names, frametime) - + + return ( + Animation(rotations, positions, orients, offsets, parents), + names, + frametime, + ) + + def load_bfa(filename, start=None, end=None, order=None, world=False): """ Reads a BVH file and constructs an animation @@ -213,31 +223,38 @@ def load_bfa(filename, start=None, end=None, order=None, world=False): names = [] orients = Quaternions.id(0) - offsets = np.array([]).reshape((0,3)) + offsets = np.array([]).reshape((0, 3)) parents = np.array([], dtype=int) for line in f: - if "HIERARCHY" in line: continue - if "MOTION" in line: continue + if "HIERARCHY" in line: + continue + if "MOTION" in line: + continue rmatch = re.match(r"ROOT (\w+)", line) if rmatch: names.append(rmatch.group(1)) - offsets = np.append(offsets, np.array([[0,0,0]]), axis=0) - orients.qs = np.append(orients.qs, np.array([[1,0,0,0]]), axis=0) - parents = np.append(parents, active) - active = (len(parents)-1) + offsets = np.append(offsets, np.array([[0, 0, 0]]), axis=0) + orients.qs = np.append(orients.qs, np.array([[1, 0, 0, 0]]), axis=0) + parents = np.append(parents, active) + active = len(parents) - 1 continue - if "{" in line: continue + if "{" in line: + continue if "}" in line: - if end_site: end_site = False - else: active = parents[active] + if end_site: + end_site = False + else: + active = parents[active] continue - offmatch = re.match(r"\s*OFFSET\s+([\-\d\.e]+)\s+([\-\d\.e]+)\s+([\-\d\.e]+)", line) + offmatch = re.match( + r"\s*OFFSET\s+([\-\d\.e]+)\s+([\-\d\.e]+)\s+([\-\d\.e]+)", line + ) if offmatch: if not end_site: offsets[active] = np.array([list(map(float, offmatch.groups()))]) @@ -255,7 +272,7 @@ def load_bfa(filename, start=None, end=None, order=None, world=False): if order is None: channelis = 0 if channels == 3 else 3 channelie = 3 if channels == 3 else 6 - parts = line.split()[2+channelis:2+channelie] + parts = line.split()[2 + channelis : 2 + channelie] if any([p not in channelmap for p in parts]): continue order = "".join([channelmap[p] for p in parts]) @@ -264,21 +281,23 @@ def load_bfa(filename, start=None, end=None, order=None, world=False): jmatch = re.match("\s*JOINT\s+(\w+)", line) if jmatch: names.append(jmatch.group(1)) - offsets = np.append(offsets, np.array([[0,0,0]]), axis=0) - orients.qs = np.append(orients.qs, np.array([[1,0,0,0]]), axis=0) - parents = np.append(parents, active) - active = (len(parents)-1) + offsets = np.append(offsets, np.array([[0, 0, 0]]), axis=0) + orients.qs = np.append(orients.qs, np.array([[1, 0, 0, 0]]), axis=0) + parents = np.append(parents, active) + active = len(parents) - 1 continue if "End Site" in line: if active + 1 in hand_idx: print("parent:", names[-1]) - name = "LeftHandIndex" if active + 1 == hand_idx[0] else "RightHandIndex" + name = ( + "LeftHandIndex" if active + 1 == hand_idx[0] else "RightHandIndex" + ) names.append(name) - offsets = np.append(offsets, np.array([[0,0,0]]), axis=0) - orients.qs = np.append(orients.qs, np.array([[1,0,0,0]]), axis=0) - parents = np.append(parents, active) - active = (len(parents)-1) + offsets = np.append(offsets, np.array([[0, 0, 0]]), axis=0) + orients.qs = np.append(orients.qs, np.array([[1, 0, 0, 0]]), axis=0) + parents = np.append(parents, active) + active = len(parents) - 1 else: end_site = True continue @@ -286,7 +305,7 @@ def load_bfa(filename, start=None, end=None, order=None, world=False): fmatch = re.match("\s*Frames:\s+(\d+)", line) if fmatch: if start and end: - fnum = (end - start)-1 + fnum = (end - start) - 1 else: fnum = int(fmatch.group(1)) jnum = len(parents) @@ -301,7 +320,7 @@ def load_bfa(filename, start=None, end=None, order=None, world=False): frametime = float(fmatch.group(1)) continue - if (start and end) and (i < start or i >= end-1): + if (start and end) and (i < start or i >= end - 1): i += 1 continue @@ -310,26 +329,31 @@ def load_bfa(filename, start=None, end=None, order=None, world=False): data_block = np.array(list(map(float, dmatch))) N = len(parents) fi = i - start if start else i - if channels == 3: + if channels == 3: # This should be root positions[0:1] & all rotations - positions[fi,0:1] = data_block[0:3] - tmp = data_block[3: ].reshape(N - 2, 3) - tmp = np.concatenate([tmp[:hand_idx[0]], - np.array([[0, 0, 0]]), - tmp[hand_idx[0]: hand_idx[1] - 1], - np.array([[0, 0, 0]]), - tmp[hand_idx[1] - 1:]], axis=0) - rotations[fi, : ] = tmp.reshape(N,3) + positions[fi, 0:1] = data_block[0:3] + tmp = data_block[3:].reshape(N - 2, 3) + tmp = np.concatenate( + [ + tmp[: hand_idx[0]], + np.array([[0, 0, 0]]), + tmp[hand_idx[0] : hand_idx[1] - 1], + np.array([[0, 0, 0]]), + tmp[hand_idx[1] - 1 :], + ], + axis=0, + ) + rotations[fi, :] = tmp.reshape(N, 3) elif channels == 6: - data_block = data_block.reshape(N,6) + data_block = data_block.reshape(N, 6) # fill in all positions - positions[fi,:] = data_block[:,0:3] - rotations[fi,:] = data_block[:,3:6] + positions[fi, :] = data_block[:, 0:3] + rotations[fi, :] = data_block[:, 3:6] elif channels == 9: - positions[fi,0] = data_block[0:3] - data_block = data_block[3:].reshape(N-1,9) - rotations[fi,1:] = data_block[:,3:6] - positions[fi,1:] += data_block[:,0:3] * data_block[:,6:9] + positions[fi, 0] = data_block[0:3] + data_block = data_block[3:].reshape(N - 1, 9) + rotations[fi, 1:] = data_block[:, 3:6] + positions[fi, 1:] += data_block[:, 0:3] * data_block[:, 6:9] else: raise Exception("Too many channels! %i" % channels) @@ -339,10 +363,22 @@ def load_bfa(filename, start=None, end=None, order=None, world=False): rotations = Quaternions.from_euler(np.radians(rotations), order=order, world=world) - return (Animation(rotations, positions, orients, offsets, parents), names, frametime) - - -def save(filename, anim, names=None, frametime=1.0/24.0, order='zyx', positions=False, orients=True): + return ( + Animation(rotations, positions, orients, offsets, parents), + names, + frametime, + ) + + +def save( + filename, + anim, + names=None, + frametime=1.0 / 24.0, + order="zyx", + positions=False, + orients=True, +): """ Saves an Animation to file as BVH @@ -373,21 +409,31 @@ def save(filename, anim, names=None, frametime=1.0/24.0, order='zyx', positions= before saving. """ - + if names is None: names = ["joint_" + str(i) for i in range(len(anim.parents))] - - with open(filename, 'w') as f: + + with open(filename, "w") as f: t = "" f.write("%sHIERARCHY\n" % t) f.write("%sROOT %s\n" % (t, names[0])) f.write("%s{\n" % t) - t += '\t' - - f.write("%sOFFSET %f %f %f\n" % (t, anim.offsets[0,0], anim.offsets[0,1], anim.offsets[0,2]) ) - f.write("%sCHANNELS 6 Xposition Yposition Zposition %s %s %s \n" % - (t, channelmap_inv[order[0]], channelmap_inv[order[1]], channelmap_inv[order[2]])) + t += "\t" + + f.write( + "%sOFFSET %f %f %f\n" + % (t, anim.offsets[0, 0], anim.offsets[0, 1], anim.offsets[0, 2]) + ) + f.write( + "%sCHANNELS 6 Xposition Yposition Zposition %s %s %s \n" + % ( + t, + channelmap_inv[order[0]], + channelmap_inv[order[1]], + channelmap_inv[order[2]], + ) + ) for i in range(anim.shape[1]): if anim.parents[i] == 0: @@ -397,64 +443,95 @@ def save(filename, anim, names=None, frametime=1.0/24.0, order='zyx', positions= f.write("%s}\n" % t) f.write("MOTION\n") - f.write("Frames: %i\n" % anim.shape[0]); - f.write("Frame Time: %f\n" % frametime); - - #if orients: + f.write("Frames: %i\n" % anim.shape[0]) + f.write("Frame Time: %f\n" % frametime) + + # if orients: # rots = np.degrees((-anim.orients[np.newaxis] * anim.rotations).euler(order=order[::-1])) - #else: + # else: # rots = np.degrees(anim.rotations.euler(order=order[::-1])) rots = np.degrees(anim.rotations.euler(order=order[::-1])) poss = anim.positions - + for i in range(anim.shape[0]): for j in range(anim.shape[1]): - + if positions or j == 0: - - f.write("%f %f %f %f %f %f " % ( - poss[i,j,0], poss[i,j,1], poss[i,j,2], - rots[i,j,ordermap[order[0]]], rots[i,j,ordermap[order[1]]], rots[i,j,ordermap[order[2]]])) - + + f.write( + "%f %f %f %f %f %f " + % ( + poss[i, j, 0], + poss[i, j, 1], + poss[i, j, 2], + rots[i, j, ordermap[order[0]]], + rots[i, j, ordermap[order[1]]], + rots[i, j, ordermap[order[2]]], + ) + ) + else: - - f.write("%f %f %f " % ( - rots[i,j,ordermap[order[0]]], rots[i,j,ordermap[order[1]]], rots[i,j,ordermap[order[2]]])) + + f.write( + "%f %f %f " + % ( + rots[i, j, ordermap[order[0]]], + rots[i, j, ordermap[order[1]]], + rots[i, j, ordermap[order[2]]], + ) + ) f.write("\n") - - -def save_joint(f, anim, names, t, i, order='zyx', positions=False): - + + +def save_joint(f, anim, names, t, i, order="zyx", positions=False): + f.write("%sJOINT %s\n" % (t, names[i])) f.write("%s{\n" % t) - t += '\t' - - f.write("%sOFFSET %f %f %f\n" % (t, anim.offsets[i,0], anim.offsets[i,1], anim.offsets[i,2])) - + t += "\t" + + f.write( + "%sOFFSET %f %f %f\n" + % (t, anim.offsets[i, 0], anim.offsets[i, 1], anim.offsets[i, 2]) + ) + if positions: - f.write("%sCHANNELS 6 Xposition Yposition Zposition %s %s %s \n" % (t, - channelmap_inv[order[0]], channelmap_inv[order[1]], channelmap_inv[order[2]])) + f.write( + "%sCHANNELS 6 Xposition Yposition Zposition %s %s %s \n" + % ( + t, + channelmap_inv[order[0]], + channelmap_inv[order[1]], + channelmap_inv[order[2]], + ) + ) else: - f.write("%sCHANNELS 3 %s %s %s\n" % (t, - channelmap_inv[order[0]], channelmap_inv[order[1]], channelmap_inv[order[2]])) - + f.write( + "%sCHANNELS 3 %s %s %s\n" + % ( + t, + channelmap_inv[order[0]], + channelmap_inv[order[1]], + channelmap_inv[order[2]], + ) + ) + end_site = True - + for j in range(anim.shape[1]): if anim.parents[j] == i: t = save_joint(f, anim, names, t, j, order=order, positions=positions) end_site = False - + if end_site: f.write("%sEnd Site\n" % t) f.write("%s{\n" % t) - t += '\t' + t += "\t" f.write("%sOFFSET %f %f %f\n" % (t, 0.0, 0.0, 0.0)) t = t[:-1] f.write("%s}\n" % t) - + t = t[:-1] f.write("%s}\n" % t) - - return t \ No newline at end of file + + return t diff --git a/utils/BVH_mod.py b/utils/BVH_mod.py index 475291a..ce7a6cb 100755 --- a/utils/BVH_mod.py +++ b/utils/BVH_mod.py @@ -4,24 +4,21 @@ from Animation import Animation from Quaternions import Quaternions -channelmap = { - 'Xrotation' : 'x', - 'Yrotation' : 'y', - 'Zrotation' : 'z' -} +channelmap = {"Xrotation": "x", "Yrotation": "y", "Zrotation": "z"} channelmap_inv = { - 'x': 'Xrotation', - 'y': 'Yrotation', - 'z': 'Zrotation', + "x": "Xrotation", + "y": "Yrotation", + "z": "Zrotation", } ordermap = { - 'x' : 0, - 'y' : 1, - 'z' : 2, + "x": 0, + "y": 1, + "z": 2, } + def load(filename, start=None, end=None, order=None, world=False, need_quater=False): """ Reads a BVH file and constructs an animation @@ -52,112 +49,119 @@ def load(filename, start=None, end=None, order=None, world=False, need_quater=Fa (animation, joint_names, frametime) Tuple of loaded animation and joint names """ - + f = open(filename, "r") i = 0 active = -1 end_site = False - + names = [] orients = Quaternions.id(0) - offsets = np.array([]).reshape((0,3)) + offsets = np.array([]).reshape((0, 3)) parents = np.array([], dtype=int) - + for line in f: - - if "HIERARCHY" in line: continue - if "MOTION" in line: continue + + if "HIERARCHY" in line: + continue + if "MOTION" in line: + continue """ Modified line read to handle mixamo data """ -# rmatch = re.match(r"ROOT (\w+)", line) + # rmatch = re.match(r"ROOT (\w+)", line) rmatch = re.match(r"ROOT (\w+:?\w+)", line) if rmatch: names.append(rmatch.group(1)) - offsets = np.append(offsets, np.array([[0,0,0]]), axis=0) - orients.qs = np.append(orients.qs, np.array([[1,0,0,0]]), axis=0) - parents = np.append(parents, active) - active = (len(parents)-1) + offsets = np.append(offsets, np.array([[0, 0, 0]]), axis=0) + orients.qs = np.append(orients.qs, np.array([[1, 0, 0, 0]]), axis=0) + parents = np.append(parents, active) + active = len(parents) - 1 continue - if "{" in line: continue + if "{" in line: + continue if "}" in line: - if end_site: end_site = False - else: active = parents[active] + if end_site: + end_site = False + else: + active = parents[active] continue - - offmatch = re.match(r"\s*OFFSET\s+([\-\d\.e]+)\s+([\-\d\.e]+)\s+([\-\d\.e]+)", line) + + offmatch = re.match( + r"\s*OFFSET\s+([\-\d\.e]+)\s+([\-\d\.e]+)\s+([\-\d\.e]+)", line + ) if offmatch: if not end_site: offsets[active] = np.array([list(map(float, offmatch.groups()))]) continue - + chanmatch = re.match(r"\s*CHANNELS\s+(\d+)", line) if chanmatch: channels = int(chanmatch.group(1)) if order is None: channelis = 0 if channels == 3 else 3 channelie = 3 if channels == 3 else 6 - parts = line.split()[2+channelis:2+channelie] + parts = line.split()[2 + channelis : 2 + channelie] if any([p not in channelmap for p in parts]): continue order = "".join([channelmap[p] for p in parts]) continue """ Modified line read to handle mixamo data """ -# jmatch = re.match("\s*JOINT\s+(\w+)", line) + # jmatch = re.match("\s*JOINT\s+(\w+)", line) jmatch = re.match("\s*JOINT\s+(\w+:?\w+)", line) if jmatch: names.append(jmatch.group(1)) - offsets = np.append(offsets, np.array([[0,0,0]]), axis=0) - orients.qs = np.append(orients.qs, np.array([[1,0,0,0]]), axis=0) - parents = np.append(parents, active) - active = (len(parents)-1) + offsets = np.append(offsets, np.array([[0, 0, 0]]), axis=0) + orients.qs = np.append(orients.qs, np.array([[1, 0, 0, 0]]), axis=0) + parents = np.append(parents, active) + active = len(parents) - 1 continue - + if "End Site" in line: end_site = True continue - + fmatch = re.match("\s*Frames:\s+(\d+)", line) if fmatch: if start and end: - fnum = (end - start)-1 + fnum = (end - start) - 1 else: fnum = int(fmatch.group(1)) jnum = len(parents) positions = offsets[np.newaxis].repeat(fnum, axis=0) rotations = np.zeros((fnum, len(orients), 3)) continue - + fmatch = re.match("\s*Frame Time:\s+([\d\.]+)", line) if fmatch: frametime = float(fmatch.group(1)) continue - - if (start and end) and (i < start or i >= end-1): + + if (start and end) and (i < start or i >= end - 1): i += 1 continue - + # dmatch = line.strip().split(' ') dmatch = line.strip().split() if dmatch: data_block = np.array(list(map(float, dmatch))) N = len(parents) fi = i - start if start else i - if channels == 3: - positions[fi,0:1] = data_block[0:3] - rotations[fi, : ] = data_block[3: ].reshape(N,3) + if channels == 3: + positions[fi, 0:1] = data_block[0:3] + rotations[fi, :] = data_block[3:].reshape(N, 3) elif channels == 6: - data_block = data_block.reshape(N,6) - positions[fi,:] = data_block[:,0:3] - rotations[fi,:] = data_block[:,3:6] + data_block = data_block.reshape(N, 6) + positions[fi, :] = data_block[:, 0:3] + rotations[fi, :] = data_block[:, 3:6] elif channels == 9: - positions[fi,0] = data_block[0:3] - data_block = data_block[3:].reshape(N-1,9) - rotations[fi,1:] = data_block[:,3:6] - positions[fi,1:] += data_block[:,0:3] * data_block[:,6:9] + positions[fi, 0] = data_block[0:3] + data_block = data_block[3:].reshape(N - 1, 9) + rotations[fi, 1:] = data_block[:, 3:6] + positions[fi, 1:] += data_block[:, 0:3] * data_block[:, 6:9] else: raise Exception("Too many channels! %i" % channels) @@ -166,16 +170,33 @@ def load(filename, start=None, end=None, order=None, world=False, need_quater=Fa f.close() if need_quater: - rotations = Quaternions.from_euler(np.radians(rotations), order=order, world=world) - elif order != 'xyz': - rotations = Quaternions.from_euler(np.radians(rotations), order=order, world=world) + rotations = Quaternions.from_euler( + np.radians(rotations), order=order, world=world + ) + elif order != "xyz": + rotations = Quaternions.from_euler( + np.radians(rotations), order=order, world=world + ) rotations = np.degrees(rotations.euler()) - - return (Animation(rotations, positions, orients, offsets, parents), names, frametime) - - -def save(filename, anim, names=None, frametime=1.0/24.0, order='zyx', positions=False, orients=True, mask=None, quater=False): + return ( + Animation(rotations, positions, orients, offsets, parents), + names, + frametime, + ) + + +def save( + filename, + anim, + names=None, + frametime=1.0 / 24.0, + order="zyx", + positions=False, + orients=True, + mask=None, + quater=False, +): """ Saves an Animation to file as BVH @@ -206,21 +227,31 @@ def save(filename, anim, names=None, frametime=1.0/24.0, order='zyx', positions= before saving. """ - + if names is None: names = ["joint_" + str(i) for i in range(len(anim.parents))] - - with open(filename, 'w') as f: + + with open(filename, "w") as f: t = "" f.write("%sHIERARCHY\n" % t) f.write("%sROOT %s\n" % (t, names[0])) f.write("%s{\n" % t) - t += '\t' + t += "\t" - f.write("%sOFFSET %f %f %f\n" % (t, anim.offsets[0,0], anim.offsets[0,1], anim.offsets[0,2]) ) - f.write("%sCHANNELS 6 Xposition Yposition Zposition %s %s %s \n" % - (t, channelmap_inv[order[0]], channelmap_inv[order[1]], channelmap_inv[order[2]])) + f.write( + "%sOFFSET %f %f %f\n" + % (t, anim.offsets[0, 0], anim.offsets[0, 1], anim.offsets[0, 2]) + ) + f.write( + "%sCHANNELS 6 Xposition Yposition Zposition %s %s %s \n" + % ( + t, + channelmap_inv[order[0]], + channelmap_inv[order[1]], + channelmap_inv[order[2]], + ) + ) for i in range(anim.shape[1]): if anim.parents[i] == 0: @@ -230,12 +261,12 @@ def save(filename, anim, names=None, frametime=1.0/24.0, order='zyx', positions= f.write("%s}\n" % t) f.write("MOTION\n") - f.write("Frames: %i\n" % anim.shape[0]); - f.write("Frame Time: %f\n" % frametime); - - #if orients: + f.write("Frames: %i\n" % anim.shape[0]) + f.write("Frame Time: %f\n" % frametime) + + # if orients: # rots = np.degrees((-anim.orients[np.newaxis] * anim.rotations).euler(order=order[::-1])) - #else: + # else: # rots = np.degrees(anim.rotations.euler(order=order[::-1])) # rots = np.degrees(anim.rotations.euler(order=order[::-1])) if quater: @@ -243,57 +274,88 @@ def save(filename, anim, names=None, frametime=1.0/24.0, order='zyx', positions= else: rots = anim.rotations poss = anim.positions - + for i in range(anim.shape[0]): for j in range(anim.shape[1]): - + if positions or j == 0: - - f.write("%f %f %f %f %f %f " % ( - poss[i,j,0], poss[i,j,1], poss[i,j,2], - rots[i,j,ordermap[order[0]]], rots[i,j,ordermap[order[1]]], rots[i,j,ordermap[order[2]]])) - + + f.write( + "%f %f %f %f %f %f " + % ( + poss[i, j, 0], + poss[i, j, 1], + poss[i, j, 2], + rots[i, j, ordermap[order[0]]], + rots[i, j, ordermap[order[1]]], + rots[i, j, ordermap[order[2]]], + ) + ) + else: if mask == None or mask[j] == 1: - f.write("%f %f %f " % ( - rots[i,j,ordermap[order[0]]], rots[i,j,ordermap[order[1]]], rots[i,j,ordermap[order[2]]])) + f.write( + "%f %f %f " + % ( + rots[i, j, ordermap[order[0]]], + rots[i, j, ordermap[order[1]]], + rots[i, j, ordermap[order[2]]], + ) + ) else: f.write("%f %f %f " % (0, 0, 0)) f.write("\n") - - -def save_joint(f, anim, names, t, i, order='zyx', positions=False): - + + +def save_joint(f, anim, names, t, i, order="zyx", positions=False): + f.write("%sJOINT %s\n" % (t, names[i])) f.write("%s{\n" % t) - t += '\t' - - f.write("%sOFFSET %f %f %f\n" % (t, anim.offsets[i,0], anim.offsets[i,1], anim.offsets[i,2])) - + t += "\t" + + f.write( + "%sOFFSET %f %f %f\n" + % (t, anim.offsets[i, 0], anim.offsets[i, 1], anim.offsets[i, 2]) + ) + if positions: - f.write("%sCHANNELS 6 Xposition Yposition Zposition %s %s %s \n" % (t, - channelmap_inv[order[0]], channelmap_inv[order[1]], channelmap_inv[order[2]])) + f.write( + "%sCHANNELS 6 Xposition Yposition Zposition %s %s %s \n" + % ( + t, + channelmap_inv[order[0]], + channelmap_inv[order[1]], + channelmap_inv[order[2]], + ) + ) else: - f.write("%sCHANNELS 3 %s %s %s\n" % (t, - channelmap_inv[order[0]], channelmap_inv[order[1]], channelmap_inv[order[2]])) - + f.write( + "%sCHANNELS 3 %s %s %s\n" + % ( + t, + channelmap_inv[order[0]], + channelmap_inv[order[1]], + channelmap_inv[order[2]], + ) + ) + end_site = True - + for j in range(anim.shape[1]): if anim.parents[j] == i: t = save_joint(f, anim, names, t, j, order=order, positions=positions) end_site = False - + if end_site: f.write("%sEnd Site\n" % t) f.write("%s{\n" % t) - t += '\t' + t += "\t" f.write("%sOFFSET %f %f %f\n" % (t, 0.0, 0.0, 0.0)) t = t[:-1] f.write("%s}\n" % t) - + t = t[:-1] f.write("%s}\n" % t) - + return t diff --git a/utils/InverseKinematics.py b/utils/InverseKinematics.py index 7e38bc9..6799c29 100755 --- a/utils/InverseKinematics.py +++ b/utils/InverseKinematics.py @@ -6,6 +6,7 @@ from Quaternions_old import Quaternions + class BasicInverseKinematics: """ Basic Inverse Kinematics Solver @@ -54,56 +55,64 @@ class BasicInverseKinematics: Optional if to suppress output defaults to False """ - + def __init__(self, animation, positions, iterations=1, silent=True): - + self.animation = animation self.positions = positions self.iterations = iterations self.silent = silent - + def __call__(self): - + children = AnimationStructure.children_list(self.animation.parents) - + for i in range(self.iterations): - + for j in AnimationStructure.joints(self.animation.parents): - + c = np.array(children[j]) - if len(c) == 0: continue - + if len(c) == 0: + continue + anim_transforms = Animation.transforms_global(self.animation) - anim_positions = anim_transforms[:,:,:3,3] + anim_positions = anim_transforms[:, :, :3, 3] anim_rotations = Quaternions.from_transforms(anim_transforms) - - jdirs = anim_positions[:,c] - anim_positions[:,np.newaxis,j] - ddirs = self.positions[:,c] - anim_positions[:,np.newaxis,j] - - jsums = np.sqrt(np.sum(jdirs**2.0, axis=-1)) + 1e-10 - dsums = np.sqrt(np.sum(ddirs**2.0, axis=-1)) + 1e-10 - - jdirs = jdirs / jsums[:,:,np.newaxis] - ddirs = ddirs / dsums[:,:,np.newaxis] - + + jdirs = anim_positions[:, c] - anim_positions[:, np.newaxis, j] + ddirs = self.positions[:, c] - anim_positions[:, np.newaxis, j] + + jsums = np.sqrt(np.sum(jdirs ** 2.0, axis=-1)) + 1e-10 + dsums = np.sqrt(np.sum(ddirs ** 2.0, axis=-1)) + 1e-10 + + jdirs = jdirs / jsums[:, :, np.newaxis] + ddirs = ddirs / dsums[:, :, np.newaxis] + angles = np.arccos(np.sum(jdirs * ddirs, axis=2).clip(-1, 1)) axises = np.cross(jdirs, ddirs) - axises = -anim_rotations[:,j,np.newaxis] * axises - + axises = -anim_rotations[:, j, np.newaxis] * axises + rotations = Quaternions.from_angle_axis(angles, axises) - + if rotations.shape[1] == 1: - averages = rotations[:,0] + averages = rotations[:, 0] else: - averages = Quaternions.exp(rotations.log().mean(axis=-2)) - - self.animation.rotations[:,j] = self.animation.rotations[:,j] * averages - + averages = Quaternions.exp(rotations.log().mean(axis=-2)) + + self.animation.rotations[:, j] = ( + self.animation.rotations[:, j] * averages + ) + if not self.silent: anim_positions = Animation.positions_global(self.animation) - error = np.mean(np.sum((anim_positions - self.positions)**2.0, axis=-1)**0.5, axis=-1) - print('[BasicInverseKinematics] Iteration %i Error: %f' % (i+1, error)) - + error = np.mean( + np.sum((anim_positions - self.positions) ** 2.0, axis=-1) ** 0.5, + axis=-1, + ) + print( + "[BasicInverseKinematics] Iteration %i Error: %f" % (i + 1, error) + ) + return self.animation @@ -161,194 +170,282 @@ class JacobianInverseKinematics: Optional if to suppress output defaults to False """ - - def __init__(self, animation, targets, - references=None, iterations=10, - recalculate=True, damping=2.0, - secondary=0.25, translate=False, - silent=False, weights=None, - weights_translate=None): - + + def __init__( + self, + animation, + targets, + references=None, + iterations=10, + recalculate=True, + damping=2.0, + secondary=0.25, + translate=False, + silent=False, + weights=None, + weights_translate=None, + ): + self.animation = animation self.targets = targets self.references = references - - self.iterations = iterations + + self.iterations = iterations self.recalculate = recalculate - self.damping = damping - self.secondary = secondary - self.translate = translate - self.silent = silent - self.weights = weights + self.damping = damping + self.secondary = secondary + self.translate = translate + self.silent = silent + self.weights = weights self.weights_translate = weights_translate - + def cross(self, a, b): o = np.empty(b.shape) - o[...,0] = a[...,1]*b[...,2] - a[...,2]*b[...,1] - o[...,1] = a[...,2]*b[...,0] - a[...,0]*b[...,2] - o[...,2] = a[...,0]*b[...,1] - a[...,1]*b[...,0] + o[..., 0] = a[..., 1] * b[..., 2] - a[..., 2] * b[..., 1] + o[..., 1] = a[..., 2] * b[..., 0] - a[..., 0] * b[..., 2] + o[..., 2] = a[..., 0] * b[..., 1] - a[..., 1] * b[..., 0] return o - + def jacobian(self, x, fp, fr, ts, dsc, tdsc): - + """ Find parent rotations """ - prs = fr[:,self.animation.parents] - prs[:,0] = Quaternions.id((1)) - + prs = fr[:, self.animation.parents] + prs[:, 0] = Quaternions.id((1)) + """ Find global positions of target joints """ - tps = fp[:,np.array(list(ts.keys()))] - + tps = fp[:, np.array(list(ts.keys()))] + """ Get partial rotations """ - qys = Quaternions.from_angle_axis(x[:,1:prs.shape[1]*3:3], np.array([[[0,1,0]]])) - qzs = Quaternions.from_angle_axis(x[:,2:prs.shape[1]*3:3], np.array([[[0,0,1]]])) - + qys = Quaternions.from_angle_axis( + x[:, 1 : prs.shape[1] * 3 : 3], np.array([[[0, 1, 0]]]) + ) + qzs = Quaternions.from_angle_axis( + x[:, 2 : prs.shape[1] * 3 : 3], np.array([[[0, 0, 1]]]) + ) + """ Find axis of rotations """ - es = np.empty((len(x),fr.shape[1]*3, 3)) - es[:,0::3] = ((prs * qzs) * qys) * np.array([[[1,0,0]]]) - es[:,1::3] = ((prs * qzs) * np.array([[[0,1,0]]])) - es[:,2::3] = ((prs * np.array([[[0,0,1]]]))) - + es = np.empty((len(x), fr.shape[1] * 3, 3)) + es[:, 0::3] = ((prs * qzs) * qys) * np.array([[[1, 0, 0]]]) + es[:, 1::3] = (prs * qzs) * np.array([[[0, 1, 0]]]) + es[:, 2::3] = prs * np.array([[[0, 0, 1]]]) + """ Construct Jacobian """ j = fp.repeat(3, axis=1) - j = dsc[np.newaxis,:,:,np.newaxis] * (tps[:,np.newaxis,:] - j[:,:,np.newaxis]) - j = self.cross(es[:,:,np.newaxis,:], j) - j = np.swapaxes(j.reshape((len(x), fr.shape[1]*3, len(ts)*3)), 1, 2) - + j = dsc[np.newaxis, :, :, np.newaxis] * ( + tps[:, np.newaxis, :] - j[:, :, np.newaxis] + ) + j = self.cross(es[:, :, np.newaxis, :], j) + j = np.swapaxes(j.reshape((len(x), fr.shape[1] * 3, len(ts) * 3)), 1, 2) + if self.translate: - - es = np.empty((len(x),fr.shape[1]*3, 3)) - es[:,0::3] = prs * np.array([[[1,0,0]]]) - es[:,1::3] = prs * np.array([[[0,1,0]]]) - es[:,2::3] = prs * np.array([[[0,0,1]]]) - - jt = tdsc[np.newaxis,:,:,np.newaxis] * es[:,:,np.newaxis,:].repeat(tps.shape[1], axis=2) - jt = np.swapaxes(jt.reshape((len(x), fr.shape[1]*3, len(ts)*3)), 1, 2) - + + es = np.empty((len(x), fr.shape[1] * 3, 3)) + es[:, 0::3] = prs * np.array([[[1, 0, 0]]]) + es[:, 1::3] = prs * np.array([[[0, 1, 0]]]) + es[:, 2::3] = prs * np.array([[[0, 0, 1]]]) + + jt = tdsc[np.newaxis, :, :, np.newaxis] * es[:, :, np.newaxis, :].repeat( + tps.shape[1], axis=2 + ) + jt = np.swapaxes(jt.reshape((len(x), fr.shape[1] * 3, len(ts) * 3)), 1, 2) + j = np.concatenate([j, jt], axis=-1) - + return j - - #@profile(immediate=True) + + # @profile(immediate=True) def __call__(self, descendants=None, gamma=1.0): - + self.descendants = descendants - + """ Calculate Masses """ if self.weights is None: self.weights = np.ones(self.animation.shape[1]) - + if self.weights_translate is None: self.weights_translate = np.ones(self.animation.shape[1]) - + """ Calculate Descendants """ if self.descendants is None: - self.descendants = AnimationStructure.descendants_mask(self.animation.parents) - + self.descendants = AnimationStructure.descendants_mask( + self.animation.parents + ) + self.tdescendants = np.eye(self.animation.shape[1]) + self.descendants - - self.first_descendants = self.descendants[:,np.array(list(self.targets.keys()))].repeat(3, axis=0).astype(int) - self.first_tdescendants = self.tdescendants[:,np.array(list(self.targets.keys()))].repeat(3, axis=0).astype(int) - + + self.first_descendants = ( + self.descendants[:, np.array(list(self.targets.keys()))] + .repeat(3, axis=0) + .astype(int) + ) + self.first_tdescendants = ( + self.tdescendants[:, np.array(list(self.targets.keys()))] + .repeat(3, axis=0) + .astype(int) + ) + """ Calculate End Effectors """ self.endeff = np.array(list(self.targets.values())) - self.endeff = np.swapaxes(self.endeff, 0, 1) - + self.endeff = np.swapaxes(self.endeff, 0, 1) + if not self.references is None: self.second_descendants = self.descendants.repeat(3, axis=0).astype(int) self.second_tdescendants = self.tdescendants.repeat(3, axis=0).astype(int) - self.second_targets = dict([(i, self.references[:,i]) for i in xrange(self.references.shape[1])]) - + self.second_targets = dict( + [(i, self.references[:, i]) for i in xrange(self.references.shape[1])] + ) + nf = len(self.animation) nj = self.animation.shape[1] - + if not self.silent: gp = Animation.positions_global(self.animation) - gp = gp[:,np.array(list(self.targets.keys()))] - error = np.mean(np.sqrt(np.sum((self.endeff - gp)**2.0, axis=2))) - print('[JacobianInverseKinematics] Start | Error: %f' % error) - + gp = gp[:, np.array(list(self.targets.keys()))] + error = np.mean(np.sqrt(np.sum((self.endeff - gp) ** 2.0, axis=2))) + print("[JacobianInverseKinematics] Start | Error: %f" % error) + for i in range(self.iterations): """ Get Global Rotations & Positions """ gt = Animation.transforms_global(self.animation) - gp = gt[:,:,:,3] - gp = gp[:,:,:3] / gp[:,:,3,np.newaxis] + gp = gt[:, :, :, 3] + gp = gp[:, :, :3] / gp[:, :, 3, np.newaxis] gr = Quaternions.from_transforms(gt) - + x = self.animation.rotations.euler().reshape(nf, -1) w = self.weights.repeat(3) - + if self.translate: x = np.hstack([x, self.animation.positions.reshape(nf, -1)]) w = np.hstack([w, self.weights_translate.repeat(3)]) - + """ Generate Jacobian """ if self.recalculate or i == 0: - j = self.jacobian(x, gp, gr, self.targets, self.first_descendants, self.first_tdescendants) - - """ Update Variables """ + j = self.jacobian( + x, + gp, + gr, + self.targets, + self.first_descendants, + self.first_tdescendants, + ) + + """ Update Variables """ l = self.damping * (1.0 / (w + 0.001)) - d = (l*l) * np.eye(x.shape[1]) - e = gamma * (self.endeff.reshape(nf,-1) - gp[:,np.array(list(self.targets.keys()))].reshape(nf, -1)) - - x += np.array(list(map(lambda jf, ef: - linalg.lu_solve(linalg.lu_factor(jf.T.dot(jf) + d), jf.T.dot(ef)), j, e))) - + d = (l * l) * np.eye(x.shape[1]) + e = gamma * ( + self.endeff.reshape(nf, -1) + - gp[:, np.array(list(self.targets.keys()))].reshape(nf, -1) + ) + + x += np.array( + list( + map( + lambda jf, ef: linalg.lu_solve( + linalg.lu_factor(jf.T.dot(jf) + d), jf.T.dot(ef) + ), + j, + e, + ) + ) + ) + """ Generate Secondary Jacobian """ if self.references is not None: - - ns = np.array(list(map(lambda jf: - np.eye(x.shape[1]) - linalg.solve(jf.T.dot(jf) + d, jf.T.dot(jf)), j))) - + + ns = np.array( + list( + map( + lambda jf: np.eye(x.shape[1]) + - linalg.solve(jf.T.dot(jf) + d, jf.T.dot(jf)), + j, + ) + ) + ) + if self.recalculate or i == 0: - j2 = self.jacobian(x, gp, gr, self.second_targets, self.second_descendants, self.second_tdescendants) - - e2 = self.secondary * (self.references.reshape(nf, -1) - gp.reshape(nf, -1)) - - x += np.array(list(map(lambda nsf, j2f, e2f: - nsf.dot(linalg.lu_solve(linalg.lu_factor(j2f.T.dot(j2f) + d), j2f.T.dot(e2f))), ns, j2, e2))) + j2 = self.jacobian( + x, + gp, + gr, + self.second_targets, + self.second_descendants, + self.second_tdescendants, + ) + + e2 = self.secondary * ( + self.references.reshape(nf, -1) - gp.reshape(nf, -1) + ) + + x += np.array( + list( + map( + lambda nsf, j2f, e2f: nsf.dot( + linalg.lu_solve( + linalg.lu_factor(j2f.T.dot(j2f) + d), j2f.T.dot(e2f) + ) + ), + ns, + j2, + e2, + ) + ) + ) """ Set Back Rotations / Translations """ self.animation.rotations = Quaternions.from_euler( - x[:,:nj*3].reshape((nf, nj, 3)), order='xyz', world=True) - + x[:, : nj * 3].reshape((nf, nj, 3)), order="xyz", world=True + ) + if self.translate: - self.animation.positions = x[:,nj*3:].reshape((nf,nj, 3)) - + self.animation.positions = x[:, nj * 3 :].reshape((nf, nj, 3)) + """ Generate Error """ - + if not self.silent: gp = Animation.positions_global(self.animation) - gp = gp[:,np.array(list(self.targets.keys()))] - error = np.mean(np.sum((self.endeff - gp)**2.0, axis=2)**0.5) - print('[JacobianInverseKinematics] Iteration %i | Error: %f' % (i+1, error)) - + gp = gp[:, np.array(list(self.targets.keys()))] + error = np.mean(np.sum((self.endeff - gp) ** 2.0, axis=2) ** 0.5) + print( + "[JacobianInverseKinematics] Iteration %i | Error: %f" + % (i + 1, error) + ) + class BasicJacobianIK: """ Same interface as BasicInverseKinematics but uses the Jacobian IK Solver Instead """ - + def __init__(self, animation, positions, iterations=10, silent=True, **kw): - - targets = dict([(i, positions[:,i]) for i in range(positions.shape[1])]) - self.ik = JacobianInverseKinematics(animation, targets, iterations=iterations, silent=silent, **kw) - + + targets = dict([(i, positions[:, i]) for i in range(positions.shape[1])]) + self.ik = JacobianInverseKinematics( + animation, targets, iterations=iterations, silent=silent, **kw + ) + def __call__(self, **kw): return self.ik(**kw) - + class ICP: - - - def __init__(self, - anim, rest, weights, mesh, goal, - find_closest=True, damping=10, - iterations=10, silent=True, - translate=True, recalculate=True, - weights_translate=None): - + def __init__( + self, + anim, + rest, + weights, + mesh, + goal, + find_closest=True, + damping=10, + iterations=10, + silent=True, + translate=True, + recalculate=True, + weights_translate=None, + ): + self.animation = anim self.rest = rest self.vweights = weights @@ -358,140 +455,170 @@ def __init__(self, self.iterations = iterations self.silent = silent self.translate = translate - self.damping = damping + self.damping = damping self.weights = None self.weights_translate = weights_translate self.recalculate = recalculate - + def cross(self, a, b): o = np.empty(b.shape) - o[...,0] = a[...,1]*b[...,2] - a[...,2]*b[...,1] - o[...,1] = a[...,2]*b[...,0] - a[...,0]*b[...,2] - o[...,2] = a[...,0]*b[...,1] - a[...,1]*b[...,0] + o[..., 0] = a[..., 1] * b[..., 2] - a[..., 2] * b[..., 1] + o[..., 1] = a[..., 2] * b[..., 0] - a[..., 0] * b[..., 2] + o[..., 2] = a[..., 0] * b[..., 1] - a[..., 1] * b[..., 0] return o - + def jacobian(self, x, fp, fr, goal, weights, des_r, des_t): - + """ Find parent rotations """ - prs = fr[:,self.animation.parents] - prs[:,0] = Quaternions.id((1)) - + prs = fr[:, self.animation.parents] + prs[:, 0] = Quaternions.id((1)) + """ Get partial rotations """ - qys = Quaternions.from_angle_axis(x[:,1:prs.shape[1]*3:3], np.array([[[0,1,0]]])) - qzs = Quaternions.from_angle_axis(x[:,2:prs.shape[1]*3:3], np.array([[[0,0,1]]])) - + qys = Quaternions.from_angle_axis( + x[:, 1 : prs.shape[1] * 3 : 3], np.array([[[0, 1, 0]]]) + ) + qzs = Quaternions.from_angle_axis( + x[:, 2 : prs.shape[1] * 3 : 3], np.array([[[0, 0, 1]]]) + ) + """ Find axis of rotations """ - es = np.empty((len(x),fr.shape[1]*3, 3)) - es[:,0::3] = ((prs * qzs) * qys) * np.array([[[1,0,0]]]) - es[:,1::3] = ((prs * qzs) * np.array([[[0,1,0]]])) - es[:,2::3] = ((prs * np.array([[[0,0,1]]]))) - + es = np.empty((len(x), fr.shape[1] * 3, 3)) + es[:, 0::3] = ((prs * qzs) * qys) * np.array([[[1, 0, 0]]]) + es[:, 1::3] = (prs * qzs) * np.array([[[0, 1, 0]]]) + es[:, 2::3] = prs * np.array([[[0, 0, 1]]]) + """ Construct Jacobian """ j = fp.repeat(3, axis=1) - j = des_r[np.newaxis,:,:,:,np.newaxis] * (goal[:,np.newaxis,:,np.newaxis] - j[:,:,np.newaxis,np.newaxis]) - j = np.sum(j * weights[np.newaxis,np.newaxis,:,:,np.newaxis], 3) - j = self.cross(es[:,:,np.newaxis,:], j) - j = np.swapaxes(j.reshape((len(x), fr.shape[1]*3, goal.shape[1]*3)), 1, 2) - + j = des_r[np.newaxis, :, :, :, np.newaxis] * ( + goal[:, np.newaxis, :, np.newaxis] - j[:, :, np.newaxis, np.newaxis] + ) + j = np.sum(j * weights[np.newaxis, np.newaxis, :, :, np.newaxis], 3) + j = self.cross(es[:, :, np.newaxis, :], j) + j = np.swapaxes(j.reshape((len(x), fr.shape[1] * 3, goal.shape[1] * 3)), 1, 2) + if self.translate: - - es = np.empty((len(x),fr.shape[1]*3, 3)) - es[:,0::3] = prs * np.array([[[1,0,0]]]) - es[:,1::3] = prs * np.array([[[0,1,0]]]) - es[:,2::3] = prs * np.array([[[0,0,1]]]) - - jt = des_t[np.newaxis,:,:,:,np.newaxis] * es[:,:,np.newaxis,np.newaxis,:].repeat(goal.shape[1], axis=2) - jt = np.sum(jt * weights[np.newaxis,np.newaxis,:,:,np.newaxis], 3) - jt = np.swapaxes(jt.reshape((len(x), fr.shape[1]*3, goal.shape[1]*3)), 1, 2) - + + es = np.empty((len(x), fr.shape[1] * 3, 3)) + es[:, 0::3] = prs * np.array([[[1, 0, 0]]]) + es[:, 1::3] = prs * np.array([[[0, 1, 0]]]) + es[:, 2::3] = prs * np.array([[[0, 0, 1]]]) + + jt = des_t[np.newaxis, :, :, :, np.newaxis] * es[ + :, :, np.newaxis, np.newaxis, : + ].repeat(goal.shape[1], axis=2) + jt = np.sum(jt * weights[np.newaxis, np.newaxis, :, :, np.newaxis], 3) + jt = np.swapaxes( + jt.reshape((len(x), fr.shape[1] * 3, goal.shape[1] * 3)), 1, 2 + ) + j = np.concatenate([j, jt], axis=-1) - + return j - - #@profile(immediate=True) + + # @profile(immediate=True) def __call__(self, descendants=None, maxjoints=4, gamma=1.0, transpose=False): - + """ Calculate Masses """ if self.weights is None: self.weights = np.ones(self.animation.shape[1]) - + if self.weights_translate is None: - self.weights_translate = np.ones(self.animation.shape[1]) - + self.weights_translate = np.ones(self.animation.shape[1]) + nf = len(self.animation) nj = self.animation.shape[1] nv = self.goal.shape[1] - - weightids = np.argsort(-self.vweights, axis=1)[:,:maxjoints] + + weightids = np.argsort(-self.vweights, axis=1)[:, :maxjoints] weightvls = np.array(list(map(lambda w, i: w[i], self.vweights, weightids))) - weightvls = weightvls / weightvls.sum(axis=1)[...,np.newaxis] - + weightvls = weightvls / weightvls.sum(axis=1)[..., np.newaxis] + if descendants is None: - self.descendants = AnimationStructure.descendants_mask(self.animation.parents) + self.descendants = AnimationStructure.descendants_mask( + self.animation.parents + ) else: self.descendants = descendants - + des_r = np.eye(nj) + self.descendants - des_r = des_r[:,weightids].repeat(3, axis=0) + des_r = des_r[:, weightids].repeat(3, axis=0) des_t = np.eye(nj) + self.descendants - des_t = des_t[:,weightids].repeat(3, axis=0) - + des_t = des_t[:, weightids].repeat(3, axis=0) + if not self.silent: - curr = Animation.skin(self.animation, self.rest, self.vweights, self.mesh, maxjoints=maxjoints) - error = np.mean(np.sqrt(np.sum((curr - self.goal)**2.0, axis=-1))) - print('[ICP] Start | Error: %f' % error) - + curr = Animation.skin( + self.animation, self.rest, self.vweights, self.mesh, maxjoints=maxjoints + ) + error = np.mean(np.sqrt(np.sum((curr - self.goal) ** 2.0, axis=-1))) + print("[ICP] Start | Error: %f" % error) + for i in range(self.iterations): - + """ Get Global Rotations & Positions """ gt = Animation.transforms_global(self.animation) - gp = gt[:,:,:,3] - gp = gp[:,:,:3] / gp[:,:,3,np.newaxis] + gp = gt[:, :, :, 3] + gp = gp[:, :, :3] / gp[:, :, 3, np.newaxis] gr = Quaternions.from_transforms(gt) - + x = self.animation.rotations.euler().reshape(nf, -1) w = self.weights.repeat(3) - + if self.translate: x = np.hstack([x, self.animation.positions.reshape(nf, -1)]) w = np.hstack([w, self.weights_translate.repeat(3)]) - + """ Get Current State """ - curr = Animation.skin(self.animation, self.rest, self.vweights, self.mesh, maxjoints=maxjoints) + curr = Animation.skin( + self.animation, self.rest, self.vweights, self.mesh, maxjoints=maxjoints + ) """ Find Cloest Points """ if self.find_closest: mapping = np.argmin( - (curr[:,:,np.newaxis] - - self.goal[:,np.newaxis,:])**2.0, axis=2) - e = gamma * (np.array(list(map(lambda g, m: g[m], self.goal, mapping))) - curr).reshape(nf, -1) + (curr[:, :, np.newaxis] - self.goal[:, np.newaxis, :]) ** 2.0, + axis=2, + ) + e = gamma * ( + np.array(list(map(lambda g, m: g[m], self.goal, mapping))) - curr + ).reshape(nf, -1) else: e = gamma * (self.goal - curr).reshape(nf, -1) - + """ Generate Jacobian """ if self.recalculate or i == 0: j = self.jacobian(x, gp, gr, self.goal, weightvls, des_r, des_t) - - """ Update Variables """ + + """ Update Variables """ l = self.damping * (1.0 / (w + 1e-10)) - d = (l*l) * np.eye(x.shape[1]) - + d = (l * l) * np.eye(x.shape[1]) + if transpose: x += np.array(list(map(lambda jf, ef: jf.T.dot(ef), j, e))) else: - x += np.array(list(map(lambda jf, ef: - linalg.lu_solve(linalg.lu_factor(jf.T.dot(jf) + d), jf.T.dot(ef)), j, e))) - + x += np.array( + list( + map( + lambda jf, ef: linalg.lu_solve( + linalg.lu_factor(jf.T.dot(jf) + d), jf.T.dot(ef) + ), + j, + e, + ) + ) + ) + """ Set Back Rotations / Translations """ self.animation.rotations = Quaternions.from_euler( - x[:,:nj*3].reshape((nf, nj, 3)), order='xyz', world=True) - + x[:, : nj * 3].reshape((nf, nj, 3)), order="xyz", world=True + ) + if self.translate: - self.animation.positions = x[:,nj*3:].reshape((nf, nj, 3)) - + self.animation.positions = x[:, nj * 3 :].reshape((nf, nj, 3)) + if not self.silent: - curr = Animation.skin(self.animation, self.rest, self.vweights, self.mesh) - error = np.mean(np.sqrt(np.sum((curr - self.goal)**2.0, axis=-1))) - print('[ICP] Iteration %i | Error: %f' % (i+1, error)) - + curr = Animation.skin( + self.animation, self.rest, self.vweights, self.mesh + ) + error = np.mean(np.sqrt(np.sum((curr - self.goal) ** 2.0, axis=-1))) + print("[ICP] Iteration %i | Error: %f" % (i + 1, error)) diff --git a/utils/Pivots.py b/utils/Pivots.py index 43c9c37..a940a22 100644 --- a/utils/Pivots.py +++ b/utils/Pivots.py @@ -2,7 +2,8 @@ from Quaternions_old import Quaternions -class Pivots: + +class Pivots: """ Pivots is an ndarray of angular rotations @@ -15,75 +16,117 @@ class Pivots: the standard arithmatic and need to be defined differently to work correctly """ - - def __init__(self, ps): self.ps = np.array(ps) - def __str__(self): return "Pivots("+ str(self.ps) + ")" - def __repr__(self): return "Pivots("+ repr(self.ps) + ")" - - def __add__(self, other): return Pivots(np.arctan2(np.sin(self.ps + other.ps), np.cos(self.ps + other.ps))) - def __sub__(self, other): return Pivots(np.arctan2(np.sin(self.ps - other.ps), np.cos(self.ps - other.ps))) - def __mul__(self, other): return Pivots(self.ps * other.ps) - def __div__(self, other): return Pivots(self.ps / other.ps) - def __mod__(self, other): return Pivots(self.ps % other.ps) - def __pow__(self, other): return Pivots(self.ps ** other.ps) - - def __lt__(self, other): return self.ps < other.ps - def __le__(self, other): return self.ps <= other.ps - def __eq__(self, other): return self.ps == other.ps - def __ne__(self, other): return self.ps != other.ps - def __ge__(self, other): return self.ps >= other.ps - def __gt__(self, other): return self.ps > other.ps - - def __abs__(self): return Pivots(abs(self.ps)) - def __neg__(self): return Pivots(-self.ps) - - def __iter__(self): return iter(self.ps) - def __len__(self): return len(self.ps) - - def __getitem__(self, k): return Pivots(self.ps[k]) - def __setitem__(self, k, v): self.ps[k] = v.ps - - def _ellipsis(self): return tuple(map(lambda x: slice(None), self.shape)) - - def quaternions(self, plane='xz'): + + def __init__(self, ps): + self.ps = np.array(ps) + + def __str__(self): + return "Pivots(" + str(self.ps) + ")" + + def __repr__(self): + return "Pivots(" + repr(self.ps) + ")" + + def __add__(self, other): + return Pivots( + np.arctan2(np.sin(self.ps + other.ps), np.cos(self.ps + other.ps)) + ) + + def __sub__(self, other): + return Pivots( + np.arctan2(np.sin(self.ps - other.ps), np.cos(self.ps - other.ps)) + ) + + def __mul__(self, other): + return Pivots(self.ps * other.ps) + + def __div__(self, other): + return Pivots(self.ps / other.ps) + + def __mod__(self, other): + return Pivots(self.ps % other.ps) + + def __pow__(self, other): + return Pivots(self.ps ** other.ps) + + def __lt__(self, other): + return self.ps < other.ps + + def __le__(self, other): + return self.ps <= other.ps + + def __eq__(self, other): + return self.ps == other.ps + + def __ne__(self, other): + return self.ps != other.ps + + def __ge__(self, other): + return self.ps >= other.ps + + def __gt__(self, other): + return self.ps > other.ps + + def __abs__(self): + return Pivots(abs(self.ps)) + + def __neg__(self): + return Pivots(-self.ps) + + def __iter__(self): + return iter(self.ps) + + def __len__(self): + return len(self.ps) + + def __getitem__(self, k): + return Pivots(self.ps[k]) + + def __setitem__(self, k, v): + self.ps[k] = v.ps + + def _ellipsis(self): + return tuple(map(lambda x: slice(None), self.shape)) + + def quaternions(self, plane="xz"): fa = self._ellipsis() axises = np.ones(self.ps.shape + (3,)) axises[fa + ("xyz".index(plane[0]),)] = 0.0 axises[fa + ("xyz".index(plane[1]),)] = 0.0 return Quaternions.from_angle_axis(self.ps, axises) - - def directions(self, plane='xz'): + + def directions(self, plane="xz"): dirs = np.zeros((len(self.ps), 3)) dirs[..., "xyz".index(plane[0])] = np.sin(self.ps) dirs[..., "xyz".index(plane[1])] = np.cos(self.ps) return dirs - + def normalized(self): xs = np.copy(self.ps) - while np.any(xs > np.pi): xs[xs > np.pi] = xs[xs > np.pi] - 2 * np.pi - while np.any(xs < -np.pi): xs[xs < -np.pi] = xs[xs < -np.pi] + 2 * np.pi + while np.any(xs > np.pi): + xs[xs > np.pi] = xs[xs > np.pi] - 2 * np.pi + while np.any(xs < -np.pi): + xs[xs < -np.pi] = xs[xs < -np.pi] + 2 * np.pi return Pivots(xs) - + def interpolate(self, ws): dir = np.average(self.directions, weights=ws, axis=0) return np.arctan2(dir[2], dir[0]) - + def copy(self): return Pivots(np.copy(self.ps)) - + @property def shape(self): return self.ps.shape - + @classmethod - def from_quaternions(cls, qs, forward='z', plane='xz'): + def from_quaternions(cls, qs, forward="z", plane="xz"): ds = np.zeros(qs.shape + (3,)) - ds[...,'xyz'.index(forward)] = 1.0 + ds[..., "xyz".index(forward)] = 1.0 return Pivots.from_directions(qs * ds, plane=plane) - + @classmethod - def from_directions(cls, ds, plane='xz'): - ys = ds[...,'xyz'.index(plane[0])] - xs = ds[...,'xyz'.index(plane[1])] + def from_directions(cls, ds, plane="xz"): + ys = ds[..., "xyz".index(plane[0])] + xs = ds[..., "xyz".index(plane[1])] return Pivots(np.arctan2(ys, xs)) - diff --git a/utils/Quaternions.py b/utils/Quaternions.py index 42991a3..b6f6896 100755 --- a/utils/Quaternions.py +++ b/utils/Quaternions.py @@ -1,5 +1,6 @@ import numpy as np + class Quaternions: """ Quaternions is a wrapper around a numpy ndarray @@ -19,10 +20,11 @@ class Quaternions: should support broadcasting and slicing in all of the usual ways. """ - + def __init__(self, qs): if isinstance(qs, np.ndarray): - if len(qs.shape) == 1: qs = np.array([qs]) + if len(qs.shape) == 1: + qs = np.array([qs]) self.qs = qs return @@ -30,42 +32,61 @@ def __init__(self, qs): self.qs = qs return - raise TypeError('Quaternions must be constructed from iterable, numpy array, or Quaternions, not %s' % type(qs)) - - def __str__(self): return "Quaternions("+ str(self.qs) + ")" - def __repr__(self): return "Quaternions("+ repr(self.qs) + ")" - + raise TypeError( + "Quaternions must be constructed from iterable, numpy array, or Quaternions, not %s" + % type(qs) + ) + + def __str__(self): + return "Quaternions(" + str(self.qs) + ")" + + def __repr__(self): + return "Quaternions(" + repr(self.qs) + ")" + """ Helper Methods for Broadcasting and Data extraction """ - + @classmethod def _broadcast(cls, sqs, oqs, scalar=False): - if isinstance(oqs, float): return sqs, oqs * np.ones(sqs.shape[:-1]) - + if isinstance(oqs, float): + return sqs, oqs * np.ones(sqs.shape[:-1]) + ss = np.array(sqs.shape) if not scalar else np.array(sqs.shape[:-1]) os = np.array(oqs.shape) if len(ss) != len(os): - raise TypeError('Quaternions cannot broadcast together shapes %s and %s' % (sqs.shape, oqs.shape)) - - if np.all(ss == os): return sqs, oqs - + raise TypeError( + "Quaternions cannot broadcast together shapes %s and %s" + % (sqs.shape, oqs.shape) + ) + + if np.all(ss == os): + return sqs, oqs + if not np.all((ss == os) | (os == np.ones(len(os))) | (ss == np.ones(len(ss)))): - raise TypeError('Quaternions cannot broadcast together shapes %s and %s' % (sqs.shape, oqs.shape)) + raise TypeError( + "Quaternions cannot broadcast together shapes %s and %s" + % (sqs.shape, oqs.shape) + ) sqsn, oqsn = sqs.copy(), oqs.copy() - for a in np.where(ss == 1)[0]: sqsn = sqsn.repeat(os[a], axis=a) - for a in np.where(os == 1)[0]: oqsn = oqsn.repeat(ss[a], axis=a) - + for a in np.where(ss == 1)[0]: + sqsn = sqsn.repeat(os[a], axis=a) + for a in np.where(os == 1)[0]: + oqsn = oqsn.repeat(ss[a], axis=a) + return sqsn, oqsn - + """ Adding Quaterions is just Defined as Multiplication """ - - def __add__(self, other): return self * other - def __sub__(self, other): return self / other - + + def __add__(self, other): + return self * other + + def __sub__(self, other): + return self / other + """ Quaterion Multiplication """ - + def __mul__(self, other): """ Quaternion multiplication has three main methods. @@ -85,36 +106,44 @@ def __mul__(self, other): Quaternions are scaled using Slerp and the identity quaternions. """ - + """ If Quaternions type do Quaternions * Quaternions """ if isinstance(other, Quaternions): sqs, oqs = Quaternions._broadcast(self.qs, other.qs) - q0 = sqs[...,0]; q1 = sqs[...,1]; - q2 = sqs[...,2]; q3 = sqs[...,3]; - r0 = oqs[...,0]; r1 = oqs[...,1]; - r2 = oqs[...,2]; r3 = oqs[...,3]; - + q0 = sqs[..., 0] + q1 = sqs[..., 1] + q2 = sqs[..., 2] + q3 = sqs[..., 3] + r0 = oqs[..., 0] + r1 = oqs[..., 1] + r2 = oqs[..., 2] + r3 = oqs[..., 3] + qs = np.empty(sqs.shape) - qs[...,0] = r0 * q0 - r1 * q1 - r2 * q2 - r3 * q3 - qs[...,1] = r0 * q1 + r1 * q0 - r2 * q3 + r3 * q2 - qs[...,2] = r0 * q2 + r1 * q3 + r2 * q0 - r3 * q1 - qs[...,3] = r0 * q3 - r1 * q2 + r2 * q1 + r3 * q0 - + qs[..., 0] = r0 * q0 - r1 * q1 - r2 * q2 - r3 * q3 + qs[..., 1] = r0 * q1 + r1 * q0 - r2 * q3 + r3 * q2 + qs[..., 2] = r0 * q2 + r1 * q3 + r2 * q0 - r3 * q1 + qs[..., 3] = r0 * q3 - r1 * q2 + r2 * q1 + r3 * q0 + return Quaternions(qs) - + """ If array type do Quaternions * Vectors """ if isinstance(other, np.ndarray) and other.shape[-1] == 3: - vs = Quaternions(np.concatenate([np.zeros(other.shape[:-1] + (1,)), other], axis=-1)) + vs = Quaternions( + np.concatenate([np.zeros(other.shape[:-1] + (1,)), other], axis=-1) + ) return (self * (vs * -self)).imaginaries """ If float do Quaternions * Scalars """ if isinstance(other, np.ndarray) or isinstance(other, float): return Quaternions.slerp(Quaternions.id_like(self), self, other) - - raise TypeError('Cannot multiply/add Quaternions with type %s' % str(type(other))) - + + raise TypeError( + "Cannot multiply/add Quaternions with type %s" % str(type(other)) + ) + def __div__(self, other): """ When a Quaternion type is supplied, division is defined @@ -124,104 +153,126 @@ def __div__(self, other): as multiplicaion of one over the supplied value. Essentially a scaling. """ - - if isinstance(other, Quaternions): return self * (-other) - if isinstance(other, np.ndarray): return self * (1.0 / other) - if isinstance(other, float): return self * (1.0 / other) - raise TypeError('Cannot divide/subtract Quaternions with type %s' + str(type(other))) - - def __eq__(self, other): return self.qs == other.qs - def __ne__(self, other): return self.qs != other.qs - + + if isinstance(other, Quaternions): + return self * (-other) + if isinstance(other, np.ndarray): + return self * (1.0 / other) + if isinstance(other, float): + return self * (1.0 / other) + raise TypeError( + "Cannot divide/subtract Quaternions with type %s" + str(type(other)) + ) + + def __eq__(self, other): + return self.qs == other.qs + + def __ne__(self, other): + return self.qs != other.qs + def __neg__(self): """ Invert Quaternions """ return Quaternions(self.qs * np.array([[1, -1, -1, -1]])) - + def __abs__(self): """ Unify Quaternions To Single Pole """ qabs = self.normalized().copy() - top = np.sum(( qabs.qs) * np.array([1,0,0,0]), axis=-1) - bot = np.sum((-qabs.qs) * np.array([1,0,0,0]), axis=-1) - qabs.qs[top < bot] = -qabs.qs[top < bot] + top = np.sum((qabs.qs) * np.array([1, 0, 0, 0]), axis=-1) + bot = np.sum((-qabs.qs) * np.array([1, 0, 0, 0]), axis=-1) + qabs.qs[top < bot] = -qabs.qs[top < bot] return qabs - - def __iter__(self): return iter(self.qs) - def __len__(self): return len(self.qs) - - def __getitem__(self, k): return Quaternions(self.qs[k]) - def __setitem__(self, k, v): self.qs[k] = v.qs - + + def __iter__(self): + return iter(self.qs) + + def __len__(self): + return len(self.qs) + + def __getitem__(self, k): + return Quaternions(self.qs[k]) + + def __setitem__(self, k, v): + self.qs[k] = v.qs + @property def lengths(self): - return np.sum(self.qs**2.0, axis=-1)**0.5 - + return np.sum(self.qs ** 2.0, axis=-1) ** 0.5 + @property def reals(self): - return self.qs[...,0] - + return self.qs[..., 0] + @property def imaginaries(self): - return self.qs[...,1:4] - + return self.qs[..., 1:4] + @property - def shape(self): return self.qs.shape[:-1] - + def shape(self): + return self.qs.shape[:-1] + def repeat(self, n, **kwargs): return Quaternions(self.qs.repeat(n, **kwargs)) - + def normalized(self): - return Quaternions(self.qs / self.lengths[...,np.newaxis]) - + return Quaternions(self.qs / self.lengths[..., np.newaxis]) + def log(self): norm = abs(self.normalized()) imgs = norm.imaginaries - lens = np.sqrt(np.sum(imgs**2, axis=-1)) + lens = np.sqrt(np.sum(imgs ** 2, axis=-1)) lens = np.arctan2(lens, norm.reals) / (lens + 1e-10) - return imgs * lens[...,np.newaxis] - + return imgs * lens[..., np.newaxis] + def constrained(self, axis): - + rl = self.reals im = np.sum(axis * self.imaginaries, axis=-1) - + t1 = -2 * np.arctan2(rl, im) + np.pi t2 = -2 * np.arctan2(rl, im) - np.pi - - top = Quaternions.exp(axis[np.newaxis] * (t1[:,np.newaxis] / 2.0)) - bot = Quaternions.exp(axis[np.newaxis] * (t2[:,np.newaxis] / 2.0)) + + top = Quaternions.exp(axis[np.newaxis] * (t1[:, np.newaxis] / 2.0)) + bot = Quaternions.exp(axis[np.newaxis] * (t2[:, np.newaxis] / 2.0)) img = self.dot(top) > self.dot(bot) - + ret = top.copy() - ret[ img] = top[ img] + ret[img] = top[img] ret[~img] = bot[~img] return ret - - def constrained_x(self): return self.constrained(np.array([1,0,0])) - def constrained_y(self): return self.constrained(np.array([0,1,0])) - def constrained_z(self): return self.constrained(np.array([0,0,1])) - - def dot(self, q): return np.sum(self.qs * q.qs, axis=-1) - - def copy(self): return Quaternions(np.copy(self.qs)) - + + def constrained_x(self): + return self.constrained(np.array([1, 0, 0])) + + def constrained_y(self): + return self.constrained(np.array([0, 1, 0])) + + def constrained_z(self): + return self.constrained(np.array([0, 0, 1])) + + def dot(self, q): + return np.sum(self.qs * q.qs, axis=-1) + + def copy(self): + return Quaternions(np.copy(self.qs)) + def reshape(self, s): self.qs.reshape(s) return self - + def interpolate(self, ws): return Quaternions.exp(np.average(abs(self).log, axis=0, weights=ws)) - - def euler(self, order='xyz'): - + + def euler(self, order="xyz"): + q = self.normalized().qs - q0 = q[...,0] - q1 = q[...,1] - q2 = q[...,2] - q3 = q[...,3] + q0 = q[..., 0] + q1 = q[..., 1] + q2 = q[..., 2] + q3 = q[..., 3] es = np.zeros(self.shape + (3,)) # These version is wrong on converting - ''' + """ if order == 'xyz': es[...,0] = np.arctan2(2 * (q0 * q1 + q2 * q3), 1 - 2 * (q1 * q1 + q2 * q2)) es[...,1] = np.arcsin((2 * (q0 * q2 - q3 * q1)).clip(-1,1)) @@ -233,18 +284,22 @@ def euler(self, order='xyz'): else: raise NotImplementedError('Cannot convert from ordering %s' % order) - ''' - - if order == 'xyz': - es[..., 2] = np.arctan2(2 * (q0 * q3 - q1 * q2), q0 * q0 + q1 * q1 - q2 * q2 - q3 * q3) - es[..., 1] = np.arcsin((2 * (q1 * q3 + q0 * q2)).clip(-1,1)) - es[..., 0] = np.arctan2(2 * (q0 * q1 - q2 * q3), q0 * q0 - q1 * q1 - q2 * q2 + q3 * q3) + """ + + if order == "xyz": + es[..., 2] = np.arctan2( + 2 * (q0 * q3 - q1 * q2), q0 * q0 + q1 * q1 - q2 * q2 - q3 * q3 + ) + es[..., 1] = np.arcsin((2 * (q1 * q3 + q0 * q2)).clip(-1, 1)) + es[..., 0] = np.arctan2( + 2 * (q0 * q1 - q2 * q3), q0 * q0 - q1 * q1 - q2 * q2 + q3 * q3 + ) else: - raise NotImplementedError('Cannot convert from ordering %s' % order) + raise NotImplementedError("Cannot convert from ordering %s" % order) # These conversion don't appear to work correctly for Maya. # http://bediyap.com/programming/convert-quaternion-to-euler-rotations/ - ''' + """ if order == 'xyz': es[..., 0] = np.arctan2(2 * (q0 * q3 - q1 * q2), q0 * q0 + q1 * q1 - q2 * q2 - q3 * q3) es[..., 1] = np.arcsin((2 * (q1 * q3 + q0 * q2)).clip(-1,1)) @@ -272,202 +327,214 @@ def euler(self, order='xyz'): else: raise KeyError('Unknown ordering %s' % order) - ''' + """ - # https://github.com/ehsan/ogre/blob/master/OgreMain/src/OgreMatrix3.cpp # Use this class and convert from matrix - + return es - - + def average(self): - + if len(self.shape) == 1: - + import numpy.core.umath_tests as ut - system = ut.matrix_multiply(self.qs[:,:,np.newaxis], self.qs[:,np.newaxis,:]).sum(axis=0) + + system = ut.matrix_multiply( + self.qs[:, :, np.newaxis], self.qs[:, np.newaxis, :] + ).sum(axis=0) w, v = np.linalg.eigh(system) - qiT_dot_qref = (self.qs[:,:,np.newaxis] * v[np.newaxis,:,:]).sum(axis=1) - return Quaternions(v[:,np.argmin((1.-qiT_dot_qref**2).sum(axis=0))]) - + qiT_dot_qref = (self.qs[:, :, np.newaxis] * v[np.newaxis, :, :]).sum(axis=1) + return Quaternions(v[:, np.argmin((1.0 - qiT_dot_qref ** 2).sum(axis=0))]) + else: - - raise NotImplementedError('Cannot average multi-dimensionsal Quaternions') + + raise NotImplementedError("Cannot average multi-dimensionsal Quaternions") def angle_axis(self): - - norm = self.normalized() - s = np.sqrt(1 - (norm.reals**2.0)) + + norm = self.normalized() + s = np.sqrt(1 - (norm.reals ** 2.0)) s[s == 0] = 0.001 - + angles = 2.0 * np.arccos(norm.reals) - axis = norm.imaginaries / s[...,np.newaxis] - + axis = norm.imaginaries / s[..., np.newaxis] + return angles, axis - - + def transforms(self): - - qw = self.qs[...,0] - qx = self.qs[...,1] - qy = self.qs[...,2] - qz = self.qs[...,3] - - x2 = qx + qx; y2 = qy + qy; z2 = qz + qz; - xx = qx * x2; yy = qy * y2; wx = qw * x2; - xy = qx * y2; yz = qy * z2; wy = qw * y2; - xz = qx * z2; zz = qz * z2; wz = qw * z2; - - m = np.empty(self.shape + (3,3)) - m[...,0,0] = 1.0 - (yy + zz) - m[...,0,1] = xy - wz - m[...,0,2] = xz + wy - m[...,1,0] = xy + wz - m[...,1,1] = 1.0 - (xx + zz) - m[...,1,2] = yz - wx - m[...,2,0] = xz - wy - m[...,2,1] = yz + wx - m[...,2,2] = 1.0 - (xx + yy) - + + qw = self.qs[..., 0] + qx = self.qs[..., 1] + qy = self.qs[..., 2] + qz = self.qs[..., 3] + + x2 = qx + qx + y2 = qy + qy + z2 = qz + qz + xx = qx * x2 + yy = qy * y2 + wx = qw * x2 + xy = qx * y2 + yz = qy * z2 + wy = qw * y2 + xz = qx * z2 + zz = qz * z2 + wz = qw * z2 + + m = np.empty(self.shape + (3, 3)) + m[..., 0, 0] = 1.0 - (yy + zz) + m[..., 0, 1] = xy - wz + m[..., 0, 2] = xz + wy + m[..., 1, 0] = xy + wz + m[..., 1, 1] = 1.0 - (xx + zz) + m[..., 1, 2] = yz - wx + m[..., 2, 0] = xz - wy + m[..., 2, 1] = yz + wx + m[..., 2, 2] = 1.0 - (xx + yy) + return m - + def ravel(self): return self.qs.ravel() - + @classmethod def id(cls, n): - + if isinstance(n, tuple): qs = np.zeros(n + (4,)) - qs[...,0] = 1.0 + qs[..., 0] = 1.0 return Quaternions(qs) - + if isinstance(n, int) or isinstance(n, long): - qs = np.zeros((n,4)) - qs[:,0] = 1.0 + qs = np.zeros((n, 4)) + qs[:, 0] = 1.0 return Quaternions(qs) - - raise TypeError('Cannot Construct Quaternion from %s type' % str(type(n))) + + raise TypeError("Cannot Construct Quaternion from %s type" % str(type(n))) @classmethod def id_like(cls, a): qs = np.zeros(a.shape + (4,)) - qs[...,0] = 1.0 + qs[..., 0] = 1.0 return Quaternions(qs) - + @classmethod def exp(cls, ws): - - ts = np.sum(ws**2.0, axis=-1)**0.5 + + ts = np.sum(ws ** 2.0, axis=-1) ** 0.5 ts[ts == 0] = 0.001 ls = np.sin(ts) / ts - + qs = np.empty(ws.shape[:-1] + (4,)) - qs[...,0] = np.cos(ts) - qs[...,1] = ws[...,0] * ls - qs[...,2] = ws[...,1] * ls - qs[...,3] = ws[...,2] * ls - + qs[..., 0] = np.cos(ts) + qs[..., 1] = ws[..., 0] * ls + qs[..., 2] = ws[..., 1] * ls + qs[..., 3] = ws[..., 2] * ls + return Quaternions(qs).normalized() - + @classmethod def slerp(cls, q0s, q1s, a): - + fst, snd = cls._broadcast(q0s.qs, q1s.qs) fst, a = cls._broadcast(fst, a, scalar=True) snd, a = cls._broadcast(snd, a, scalar=True) - + len = np.sum(fst * snd, axis=-1) - + neg = len < 0.0 len[neg] = -len[neg] snd[neg] = -snd[neg] - + amount0 = np.zeros(a.shape) amount1 = np.zeros(a.shape) linear = (1.0 - len) < 0.01 omegas = np.arccos(len[~linear]) sinoms = np.sin(omegas) - - amount0[ linear] = 1.0 - a[linear] - amount1[ linear] = a[linear] + + amount0[linear] = 1.0 - a[linear] + amount1[linear] = a[linear] amount0[~linear] = np.sin((1.0 - a[~linear]) * omegas) / sinoms - amount1[~linear] = np.sin( a[~linear] * omegas) / sinoms - + amount1[~linear] = np.sin(a[~linear] * omegas) / sinoms + return Quaternions( - amount0[...,np.newaxis] * fst + - amount1[...,np.newaxis] * snd) - + amount0[..., np.newaxis] * fst + amount1[..., np.newaxis] * snd + ) + @classmethod def between(cls, v0s, v1s): a = np.cross(v0s, v1s) - w = np.sqrt((v0s**2).sum(axis=-1) * (v1s**2).sum(axis=-1)) + (v0s * v1s).sum(axis=-1) - return Quaternions(np.concatenate([w[...,np.newaxis], a], axis=-1)).normalized() - + w = np.sqrt((v0s ** 2).sum(axis=-1) * (v1s ** 2).sum(axis=-1)) + ( + v0s * v1s + ).sum(axis=-1) + return Quaternions( + np.concatenate([w[..., np.newaxis], a], axis=-1) + ).normalized() + @classmethod def from_angle_axis(cls, angles, axis): - axis = axis / (np.sqrt(np.sum(axis**2, axis=-1)) + 1e-10)[...,np.newaxis] - sines = np.sin(angles / 2.0)[...,np.newaxis] - cosines = np.cos(angles / 2.0)[...,np.newaxis] + axis = axis / (np.sqrt(np.sum(axis ** 2, axis=-1)) + 1e-10)[..., np.newaxis] + sines = np.sin(angles / 2.0)[..., np.newaxis] + cosines = np.cos(angles / 2.0)[..., np.newaxis] return Quaternions(np.concatenate([cosines, axis * sines], axis=-1)) - + @classmethod - def from_euler(cls, es, order='xyz', world=False): - + def from_euler(cls, es, order="xyz", world=False): + axis = { - 'x' : np.array([1,0,0]), - 'y' : np.array([0,1,0]), - 'z' : np.array([0,0,1]), + "x": np.array([1, 0, 0]), + "y": np.array([0, 1, 0]), + "z": np.array([0, 0, 1]), } - - q0s = Quaternions.from_angle_axis(es[...,0], axis[order[0]]) - q1s = Quaternions.from_angle_axis(es[...,1], axis[order[1]]) - q2s = Quaternions.from_angle_axis(es[...,2], axis[order[2]]) - + + q0s = Quaternions.from_angle_axis(es[..., 0], axis[order[0]]) + q1s = Quaternions.from_angle_axis(es[..., 1], axis[order[1]]) + q2s = Quaternions.from_angle_axis(es[..., 2], axis[order[2]]) + return (q2s * (q1s * q0s)) if world else (q0s * (q1s * q2s)) - + @classmethod def from_transforms(cls, ts): - - d0, d1, d2 = ts[...,0,0], ts[...,1,1], ts[...,2,2] - - q0 = ( d0 + d1 + d2 + 1.0) / 4.0 - q1 = ( d0 - d1 - d2 + 1.0) / 4.0 + + d0, d1, d2 = ts[..., 0, 0], ts[..., 1, 1], ts[..., 2, 2] + + q0 = (d0 + d1 + d2 + 1.0) / 4.0 + q1 = (d0 - d1 - d2 + 1.0) / 4.0 q2 = (-d0 + d1 - d2 + 1.0) / 4.0 q3 = (-d0 - d1 + d2 + 1.0) / 4.0 - - q0 = np.sqrt(q0.clip(0,None)) - q1 = np.sqrt(q1.clip(0,None)) - q2 = np.sqrt(q2.clip(0,None)) - q3 = np.sqrt(q3.clip(0,None)) - + + q0 = np.sqrt(q0.clip(0, None)) + q1 = np.sqrt(q1.clip(0, None)) + q2 = np.sqrt(q2.clip(0, None)) + q3 = np.sqrt(q3.clip(0, None)) + c0 = (q0 >= q1) & (q0 >= q2) & (q0 >= q3) c1 = (q1 >= q0) & (q1 >= q2) & (q1 >= q3) c2 = (q2 >= q0) & (q2 >= q1) & (q2 >= q3) c3 = (q3 >= q0) & (q3 >= q1) & (q3 >= q2) - - q1[c0] *= np.sign(ts[c0,2,1] - ts[c0,1,2]) - q2[c0] *= np.sign(ts[c0,0,2] - ts[c0,2,0]) - q3[c0] *= np.sign(ts[c0,1,0] - ts[c0,0,1]) - - q0[c1] *= np.sign(ts[c1,2,1] - ts[c1,1,2]) - q2[c1] *= np.sign(ts[c1,1,0] + ts[c1,0,1]) - q3[c1] *= np.sign(ts[c1,0,2] + ts[c1,2,0]) - - q0[c2] *= np.sign(ts[c2,0,2] - ts[c2,2,0]) - q1[c2] *= np.sign(ts[c2,1,0] + ts[c2,0,1]) - q3[c2] *= np.sign(ts[c2,2,1] + ts[c2,1,2]) - - q0[c3] *= np.sign(ts[c3,1,0] - ts[c3,0,1]) - q1[c3] *= np.sign(ts[c3,2,0] + ts[c3,0,2]) - q2[c3] *= np.sign(ts[c3,2,1] + ts[c3,1,2]) - + + q1[c0] *= np.sign(ts[c0, 2, 1] - ts[c0, 1, 2]) + q2[c0] *= np.sign(ts[c0, 0, 2] - ts[c0, 2, 0]) + q3[c0] *= np.sign(ts[c0, 1, 0] - ts[c0, 0, 1]) + + q0[c1] *= np.sign(ts[c1, 2, 1] - ts[c1, 1, 2]) + q2[c1] *= np.sign(ts[c1, 1, 0] + ts[c1, 0, 1]) + q3[c1] *= np.sign(ts[c1, 0, 2] + ts[c1, 2, 0]) + + q0[c2] *= np.sign(ts[c2, 0, 2] - ts[c2, 2, 0]) + q1[c2] *= np.sign(ts[c2, 1, 0] + ts[c2, 0, 1]) + q3[c2] *= np.sign(ts[c2, 2, 1] + ts[c2, 1, 2]) + + q0[c3] *= np.sign(ts[c3, 1, 0] - ts[c3, 0, 1]) + q1[c3] *= np.sign(ts[c3, 2, 0] + ts[c3, 0, 2]) + q2[c3] *= np.sign(ts[c3, 2, 1] + ts[c3, 1, 2]) + qs = np.empty(ts.shape[:-2] + (4,)) - qs[...,0] = q0 - qs[...,1] = q1 - qs[...,2] = q2 - qs[...,3] = q3 - + qs[..., 0] = q0 + qs[..., 1] = q1 + qs[..., 2] = q2 + qs[..., 3] = q3 + return cls(qs) diff --git a/utils/Quaternions_old.py b/utils/Quaternions_old.py index fa43129..74afa31 100755 --- a/utils/Quaternions_old.py +++ b/utils/Quaternions_old.py @@ -1,5 +1,6 @@ import numpy as np + class Quaternions: """ Quaternions is a wrapper around a numpy ndarray @@ -19,54 +20,74 @@ class Quaternions: should support broadcasting and slicing in all of the usual ways. """ - + def __init__(self, qs): if isinstance(qs, np.ndarray): - - if len(qs.shape) == 1: qs = np.array([qs]) + + if len(qs.shape) == 1: + qs = np.array([qs]) self.qs = qs return - + if isinstance(qs, Quaternions): self.qs = qs.qs return - - raise TypeError('Quaternions must be constructed from iterable, numpy array, or Quaternions, not %s' % type(qs)) - - def __str__(self): return "Quaternions("+ str(self.qs) + ")" - def __repr__(self): return "Quaternions("+ repr(self.qs) + ")" - + + raise TypeError( + "Quaternions must be constructed from iterable, numpy array, or Quaternions, not %s" + % type(qs) + ) + + def __str__(self): + return "Quaternions(" + str(self.qs) + ")" + + def __repr__(self): + return "Quaternions(" + repr(self.qs) + ")" + """ Helper Methods for Broadcasting and Data extraction """ - + @classmethod def _broadcast(cls, sqs, oqs, scalar=False): - if isinstance(oqs, float): return sqs, oqs * np.ones(sqs.shape[:-1]) - + if isinstance(oqs, float): + return sqs, oqs * np.ones(sqs.shape[:-1]) + ss = np.array(sqs.shape) if not scalar else np.array(sqs.shape[:-1]) os = np.array(oqs.shape) if len(ss) != len(os): - raise TypeError('Quaternions cannot broadcast together shapes %s and %s' % (sqs.shape, oqs.shape)) - - if np.all(ss == os): return sqs, oqs - + raise TypeError( + "Quaternions cannot broadcast together shapes %s and %s" + % (sqs.shape, oqs.shape) + ) + + if np.all(ss == os): + return sqs, oqs + if not np.all((ss == os) | (os == np.ones(len(os))) | (ss == np.ones(len(ss)))): - raise TypeError('Quaternions cannot broadcast together shapes %s and %s' % (sqs.shape, oqs.shape)) + raise TypeError( + "Quaternions cannot broadcast together shapes %s and %s" + % (sqs.shape, oqs.shape) + ) sqsn, oqsn = sqs.copy(), oqs.copy() - for a in np.where(ss == 1)[0]: sqsn = sqsn.repeat(os[a], axis=a) - for a in np.where(os == 1)[0]: oqsn = oqsn.repeat(ss[a], axis=a) - + for a in np.where(ss == 1)[0]: + sqsn = sqsn.repeat(os[a], axis=a) + for a in np.where(os == 1)[0]: + oqsn = oqsn.repeat(ss[a], axis=a) + return sqsn, oqsn - + """ Adding Quaterions is just Defined as Multiplication """ - - def __add__(self, other): return self * other - def __sub__(self, other): return self / other - + + def __add__(self, other): + return self * other + + def __sub__(self, other): + return self / other + """ Quaterion Multiplication """ - + def __mul__(self, other): """ Quaternion multiplication has three main methods. @@ -86,36 +107,44 @@ def __mul__(self, other): Quaternions are scaled using Slerp and the identity quaternions. """ - + """ If Quaternions type do Quaternions * Quaternions """ if isinstance(other, Quaternions): sqs, oqs = Quaternions._broadcast(self.qs, other.qs) - q0 = sqs[...,0]; q1 = sqs[...,1]; - q2 = sqs[...,2]; q3 = sqs[...,3]; - r0 = oqs[...,0]; r1 = oqs[...,1]; - r2 = oqs[...,2]; r3 = oqs[...,3]; - + q0 = sqs[..., 0] + q1 = sqs[..., 1] + q2 = sqs[..., 2] + q3 = sqs[..., 3] + r0 = oqs[..., 0] + r1 = oqs[..., 1] + r2 = oqs[..., 2] + r3 = oqs[..., 3] + qs = np.empty(sqs.shape) - qs[...,0] = r0 * q0 - r1 * q1 - r2 * q2 - r3 * q3 - qs[...,1] = r0 * q1 + r1 * q0 - r2 * q3 + r3 * q2 - qs[...,2] = r0 * q2 + r1 * q3 + r2 * q0 - r3 * q1 - qs[...,3] = r0 * q3 - r1 * q2 + r2 * q1 + r3 * q0 - + qs[..., 0] = r0 * q0 - r1 * q1 - r2 * q2 - r3 * q3 + qs[..., 1] = r0 * q1 + r1 * q0 - r2 * q3 + r3 * q2 + qs[..., 2] = r0 * q2 + r1 * q3 + r2 * q0 - r3 * q1 + qs[..., 3] = r0 * q3 - r1 * q2 + r2 * q1 + r3 * q0 + return Quaternions(qs) - + """ If array type do Quaternions * Vectors """ if isinstance(other, np.ndarray) and other.shape[-1] == 3: - vs = Quaternions(np.concatenate([np.zeros(other.shape[:-1] + (1,)), other], axis=-1)) + vs = Quaternions( + np.concatenate([np.zeros(other.shape[:-1] + (1,)), other], axis=-1) + ) return (self * (vs * -self)).imaginaries - + """ If float do Quaternions * Scalars """ if isinstance(other, np.ndarray) or isinstance(other, float): return Quaternions.slerp(Quaternions.id_like(self), self, other) - - raise TypeError('Cannot multiply/add Quaternions with type %s' % str(type(other))) - + + raise TypeError( + "Cannot multiply/add Quaternions with type %s" % str(type(other)) + ) + def __div__(self, other): """ When a Quaternion type is supplied, division is defined @@ -125,113 +154,143 @@ def __div__(self, other): as multiplicaion of one over the supplied value. Essentially a scaling. """ - - if isinstance(other, Quaternions): return self * (-other) - if isinstance(other, np.ndarray): return self * (1.0 / other) - if isinstance(other, float): return self * (1.0 / other) - raise TypeError('Cannot divide/subtract Quaternions with type %s' + str(type(other))) - - def __eq__(self, other): return self.qs == other.qs - def __ne__(self, other): return self.qs != other.qs - + + if isinstance(other, Quaternions): + return self * (-other) + if isinstance(other, np.ndarray): + return self * (1.0 / other) + if isinstance(other, float): + return self * (1.0 / other) + raise TypeError( + "Cannot divide/subtract Quaternions with type %s" + str(type(other)) + ) + + def __eq__(self, other): + return self.qs == other.qs + + def __ne__(self, other): + return self.qs != other.qs + def __neg__(self): """ Invert Quaternions """ return Quaternions(self.qs * np.array([[1, -1, -1, -1]])) - + def __abs__(self): """ Unify Quaternions To Single Pole """ qabs = self.normalized().copy() - top = np.sum(( qabs.qs) * np.array([1,0,0,0]), axis=-1) - bot = np.sum((-qabs.qs) * np.array([1,0,0,0]), axis=-1) - qabs.qs[top < bot] = -qabs.qs[top < bot] + top = np.sum((qabs.qs) * np.array([1, 0, 0, 0]), axis=-1) + bot = np.sum((-qabs.qs) * np.array([1, 0, 0, 0]), axis=-1) + qabs.qs[top < bot] = -qabs.qs[top < bot] return qabs - - def __iter__(self): return iter(self.qs) - def __len__(self): return len(self.qs) - - def __getitem__(self, k): return Quaternions(self.qs[k]) - def __setitem__(self, k, v): self.qs[k] = v.qs - + + def __iter__(self): + return iter(self.qs) + + def __len__(self): + return len(self.qs) + + def __getitem__(self, k): + return Quaternions(self.qs[k]) + + def __setitem__(self, k, v): + self.qs[k] = v.qs + @property def lengths(self): - return np.sum(self.qs**2.0, axis=-1)**0.5 - + return np.sum(self.qs ** 2.0, axis=-1) ** 0.5 + @property def reals(self): - return self.qs[...,0] - + return self.qs[..., 0] + @property def imaginaries(self): - return self.qs[...,1:4] - + return self.qs[..., 1:4] + @property - def shape(self): return self.qs.shape[:-1] - + def shape(self): + return self.qs.shape[:-1] + def repeat(self, n, **kwargs): return Quaternions(self.qs.repeat(n, **kwargs)) - + def normalized(self): - return Quaternions(self.qs / self.lengths[...,np.newaxis]) - + return Quaternions(self.qs / self.lengths[..., np.newaxis]) + def log(self): norm = abs(self.normalized()) imgs = norm.imaginaries - lens = np.sqrt(np.sum(imgs**2, axis=-1)) + lens = np.sqrt(np.sum(imgs ** 2, axis=-1)) lens = np.arctan2(lens, norm.reals) / (lens + 1e-10) - return imgs * lens[...,np.newaxis] - + return imgs * lens[..., np.newaxis] + def constrained(self, axis): - + rl = self.reals im = np.sum(axis * self.imaginaries, axis=-1) - + t1 = -2 * np.arctan2(rl, im) + np.pi t2 = -2 * np.arctan2(rl, im) - np.pi - - top = Quaternions.exp(axis[np.newaxis] * (t1[:,np.newaxis] / 2.0)) - bot = Quaternions.exp(axis[np.newaxis] * (t2[:,np.newaxis] / 2.0)) + + top = Quaternions.exp(axis[np.newaxis] * (t1[:, np.newaxis] / 2.0)) + bot = Quaternions.exp(axis[np.newaxis] * (t2[:, np.newaxis] / 2.0)) img = self.dot(top) > self.dot(bot) - + ret = top.copy() - ret[ img] = top[ img] + ret[img] = top[img] ret[~img] = bot[~img] return ret - - def constrained_x(self): return self.constrained(np.array([1,0,0])) - def constrained_y(self): return self.constrained(np.array([0,1,0])) - def constrained_z(self): return self.constrained(np.array([0,0,1])) - - def dot(self, q): return np.sum(self.qs * q.qs, axis=-1) - - def copy(self): return Quaternions(np.copy(self.qs)) - + + def constrained_x(self): + return self.constrained(np.array([1, 0, 0])) + + def constrained_y(self): + return self.constrained(np.array([0, 1, 0])) + + def constrained_z(self): + return self.constrained(np.array([0, 0, 1])) + + def dot(self, q): + return np.sum(self.qs * q.qs, axis=-1) + + def copy(self): + return Quaternions(np.copy(self.qs)) + def reshape(self, s): self.qs.reshape(s) return self - + def interpolate(self, ws): return Quaternions.exp(np.average(abs(self).log, axis=0, weights=ws)) - - def euler(self, order='xyz'): - + + def euler(self, order="xyz"): + q = self.normalized().qs - q0 = q[...,0] - q1 = q[...,1] - q2 = q[...,2] - q3 = q[...,3] + q0 = q[..., 0] + q1 = q[..., 1] + q2 = q[..., 2] + q3 = q[..., 3] es = np.zeros(self.shape + (3,)) - - if order == 'xyz': - es[...,0] = np.arctan2(2 * (q0 * q1 + q2 * q3), 1 - 2 * (q1 * q1 + q2 * q2)) - es[...,1] = np.arcsin((2 * (q0 * q2 - q3 * q1)).clip(-1,1)) - es[...,2] = np.arctan2(2 * (q0 * q3 + q1 * q2), 1 - 2 * (q2 * q2 + q3 * q3)) - elif order == 'yzx': - es[...,0] = np.arctan2(2 * (q1 * q0 - q2 * q3), -q1 * q1 + q2 * q2 - q3 * q3 + q0 * q0) - es[...,1] = np.arctan2(2 * (q2 * q0 - q1 * q3), q1 * q1 - q2 * q2 - q3 * q3 + q0 * q0) - es[...,2] = np.arcsin((2 * (q1 * q2 + q3 * q0)).clip(-1,1)) + + if order == "xyz": + es[..., 0] = np.arctan2( + 2 * (q0 * q1 + q2 * q3), 1 - 2 * (q1 * q1 + q2 * q2) + ) + es[..., 1] = np.arcsin((2 * (q0 * q2 - q3 * q1)).clip(-1, 1)) + es[..., 2] = np.arctan2( + 2 * (q0 * q3 + q1 * q2), 1 - 2 * (q2 * q2 + q3 * q3) + ) + elif order == "yzx": + es[..., 0] = np.arctan2( + 2 * (q1 * q0 - q2 * q3), -q1 * q1 + q2 * q2 - q3 * q3 + q0 * q0 + ) + es[..., 1] = np.arctan2( + 2 * (q2 * q0 - q1 * q3), q1 * q1 - q2 * q2 - q3 * q3 + q0 * q0 + ) + es[..., 2] = np.arcsin((2 * (q1 * q2 + q3 * q0)).clip(-1, 1)) else: - raise NotImplementedError('Cannot convert from ordering %s' % order) - + raise NotImplementedError("Cannot convert from ordering %s" % order) + """ # These conversion don't appear to work correctly for Maya. @@ -265,203 +324,213 @@ def euler(self, order='xyz'): raise KeyError('Unknown ordering %s' % order) """ - + # https://github.com/ehsan/ogre/blob/master/OgreMain/src/OgreMatrix3.cpp # Use this class and convert from matrix - + return es - - + def average(self): - + if len(self.shape) == 1: - + import numpy.core.umath_tests as ut - system = ut.matrix_multiply(self.qs[:,:,np.newaxis], self.qs[:,np.newaxis,:]).sum(axis=0) + + system = ut.matrix_multiply( + self.qs[:, :, np.newaxis], self.qs[:, np.newaxis, :] + ).sum(axis=0) w, v = np.linalg.eigh(system) - qiT_dot_qref = (self.qs[:,:,np.newaxis] * v[np.newaxis,:,:]).sum(axis=1) - return Quaternions(v[:,np.argmin((1.-qiT_dot_qref**2).sum(axis=0))]) - + qiT_dot_qref = (self.qs[:, :, np.newaxis] * v[np.newaxis, :, :]).sum(axis=1) + return Quaternions(v[:, np.argmin((1.0 - qiT_dot_qref ** 2).sum(axis=0))]) + else: - - raise NotImplementedError('Cannot average multi-dimensionsal Quaternions') + + raise NotImplementedError("Cannot average multi-dimensionsal Quaternions") def angle_axis(self): - - norm = self.normalized() - s = np.sqrt(1 - (norm.reals**2.0)) + + norm = self.normalized() + s = np.sqrt(1 - (norm.reals ** 2.0)) s[s == 0] = 0.001 - + angles = 2.0 * np.arccos(norm.reals) - axis = norm.imaginaries / s[...,np.newaxis] - + axis = norm.imaginaries / s[..., np.newaxis] + return angles, axis - - + def transforms(self): - - qw = self.qs[...,0] - qx = self.qs[...,1] - qy = self.qs[...,2] - qz = self.qs[...,3] - - x2 = qx + qx; y2 = qy + qy; z2 = qz + qz; - xx = qx * x2; yy = qy * y2; wx = qw * x2; - xy = qx * y2; yz = qy * z2; wy = qw * y2; - xz = qx * z2; zz = qz * z2; wz = qw * z2; - - m = np.empty(self.shape + (3,3)) - m[...,0,0] = 1.0 - (yy + zz) - m[...,0,1] = xy - wz - m[...,0,2] = xz + wy - m[...,1,0] = xy + wz - m[...,1,1] = 1.0 - (xx + zz) - m[...,1,2] = yz - wx - m[...,2,0] = xz - wy - m[...,2,1] = yz + wx - m[...,2,2] = 1.0 - (xx + yy) - + + qw = self.qs[..., 0] + qx = self.qs[..., 1] + qy = self.qs[..., 2] + qz = self.qs[..., 3] + + x2 = qx + qx + y2 = qy + qy + z2 = qz + qz + xx = qx * x2 + yy = qy * y2 + wx = qw * x2 + xy = qx * y2 + yz = qy * z2 + wy = qw * y2 + xz = qx * z2 + zz = qz * z2 + wz = qw * z2 + + m = np.empty(self.shape + (3, 3)) + m[..., 0, 0] = 1.0 - (yy + zz) + m[..., 0, 1] = xy - wz + m[..., 0, 2] = xz + wy + m[..., 1, 0] = xy + wz + m[..., 1, 1] = 1.0 - (xx + zz) + m[..., 1, 2] = yz - wx + m[..., 2, 0] = xz - wy + m[..., 2, 1] = yz + wx + m[..., 2, 2] = 1.0 - (xx + yy) + return m - + def ravel(self): return self.qs.ravel() - + @classmethod def id(cls, n): - + if isinstance(n, tuple): qs = np.zeros(n + (4,)) - qs[...,0] = 1.0 + qs[..., 0] = 1.0 return Quaternions(qs) - + if isinstance(n, int) or isinstance(n, long): - qs = np.zeros((n,4)) - qs[:,0] = 1.0 + qs = np.zeros((n, 4)) + qs[:, 0] = 1.0 return Quaternions(qs) - - raise TypeError('Cannot Construct Quaternion from %s type' % str(type(n))) + + raise TypeError("Cannot Construct Quaternion from %s type" % str(type(n))) @classmethod def id_like(cls, a): qs = np.zeros(a.shape + (4,)) - qs[...,0] = 1.0 + qs[..., 0] = 1.0 return Quaternions(qs) - + @classmethod def exp(cls, ws): - - ts = np.sum(ws**2.0, axis=-1)**0.5 + + ts = np.sum(ws ** 2.0, axis=-1) ** 0.5 ts[ts == 0] = 0.001 ls = np.sin(ts) / ts - + qs = np.empty(ws.shape[:-1] + (4,)) - qs[...,0] = np.cos(ts) - qs[...,1] = ws[...,0] * ls - qs[...,2] = ws[...,1] * ls - qs[...,3] = ws[...,2] * ls - + qs[..., 0] = np.cos(ts) + qs[..., 1] = ws[..., 0] * ls + qs[..., 2] = ws[..., 1] * ls + qs[..., 3] = ws[..., 2] * ls + return Quaternions(qs).normalized() - + @classmethod def slerp(cls, q0s, q1s, a): - + fst, snd = cls._broadcast(q0s.qs, q1s.qs) fst, a = cls._broadcast(fst, a, scalar=True) snd, a = cls._broadcast(snd, a, scalar=True) - + len = np.sum(fst * snd, axis=-1) - + neg = len < 0.0 len[neg] = -len[neg] snd[neg] = -snd[neg] - + amount0 = np.zeros(a.shape) amount1 = np.zeros(a.shape) linear = (1.0 - len) < 0.01 omegas = np.arccos(len[~linear]) sinoms = np.sin(omegas) - - amount0[ linear] = 1.0 - a[linear] - amount1[ linear] = a[linear] + + amount0[linear] = 1.0 - a[linear] + amount1[linear] = a[linear] amount0[~linear] = np.sin((1.0 - a[~linear]) * omegas) / sinoms - amount1[~linear] = np.sin( a[~linear] * omegas) / sinoms - + amount1[~linear] = np.sin(a[~linear] * omegas) / sinoms + return Quaternions( - amount0[...,np.newaxis] * fst + - amount1[...,np.newaxis] * snd) - + amount0[..., np.newaxis] * fst + amount1[..., np.newaxis] * snd + ) + @classmethod def between(cls, v0s, v1s): a = np.cross(v0s, v1s) - w = np.sqrt((v0s**2).sum(axis=-1) * (v1s**2).sum(axis=-1)) + (v0s * v1s).sum(axis=-1) - return Quaternions(np.concatenate([w[...,np.newaxis], a], axis=-1)).normalized() - + w = np.sqrt((v0s ** 2).sum(axis=-1) * (v1s ** 2).sum(axis=-1)) + ( + v0s * v1s + ).sum(axis=-1) + return Quaternions( + np.concatenate([w[..., np.newaxis], a], axis=-1) + ).normalized() + @classmethod def from_angle_axis(cls, angles, axis): - axis = axis / (np.sqrt(np.sum(axis**2, axis=-1)) + 1e-10)[...,np.newaxis] - sines = np.sin(angles / 2.0)[...,np.newaxis] - cosines = np.cos(angles / 2.0)[...,np.newaxis] + axis = axis / (np.sqrt(np.sum(axis ** 2, axis=-1)) + 1e-10)[..., np.newaxis] + sines = np.sin(angles / 2.0)[..., np.newaxis] + cosines = np.cos(angles / 2.0)[..., np.newaxis] return Quaternions(np.concatenate([cosines, axis * sines], axis=-1)) - + @classmethod - def from_euler(cls, es, order='xyz', world=False): - + def from_euler(cls, es, order="xyz", world=False): + axis = { - 'x' : np.array([1,0,0]), - 'y' : np.array([0,1,0]), - 'z' : np.array([0,0,1]), + "x": np.array([1, 0, 0]), + "y": np.array([0, 1, 0]), + "z": np.array([0, 0, 1]), } - - q0s = Quaternions.from_angle_axis(es[...,0], axis[order[0]]) - q1s = Quaternions.from_angle_axis(es[...,1], axis[order[1]]) - q2s = Quaternions.from_angle_axis(es[...,2], axis[order[2]]) - + + q0s = Quaternions.from_angle_axis(es[..., 0], axis[order[0]]) + q1s = Quaternions.from_angle_axis(es[..., 1], axis[order[1]]) + q2s = Quaternions.from_angle_axis(es[..., 2], axis[order[2]]) + return (q2s * (q1s * q0s)) if world else (q0s * (q1s * q2s)) - + @classmethod def from_transforms(cls, ts): - - d0, d1, d2 = ts[...,0,0], ts[...,1,1], ts[...,2,2] - - q0 = ( d0 + d1 + d2 + 1.0) / 4.0 - q1 = ( d0 - d1 - d2 + 1.0) / 4.0 + + d0, d1, d2 = ts[..., 0, 0], ts[..., 1, 1], ts[..., 2, 2] + + q0 = (d0 + d1 + d2 + 1.0) / 4.0 + q1 = (d0 - d1 - d2 + 1.0) / 4.0 q2 = (-d0 + d1 - d2 + 1.0) / 4.0 q3 = (-d0 - d1 + d2 + 1.0) / 4.0 - - q0 = np.sqrt(q0.clip(0,None)) - q1 = np.sqrt(q1.clip(0,None)) - q2 = np.sqrt(q2.clip(0,None)) - q3 = np.sqrt(q3.clip(0,None)) - + + q0 = np.sqrt(q0.clip(0, None)) + q1 = np.sqrt(q1.clip(0, None)) + q2 = np.sqrt(q2.clip(0, None)) + q3 = np.sqrt(q3.clip(0, None)) + c0 = (q0 >= q1) & (q0 >= q2) & (q0 >= q3) c1 = (q1 >= q0) & (q1 >= q2) & (q1 >= q3) c2 = (q2 >= q0) & (q2 >= q1) & (q2 >= q3) c3 = (q3 >= q0) & (q3 >= q1) & (q3 >= q2) - - q1[c0] *= np.sign(ts[c0,2,1] - ts[c0,1,2]) - q2[c0] *= np.sign(ts[c0,0,2] - ts[c0,2,0]) - q3[c0] *= np.sign(ts[c0,1,0] - ts[c0,0,1]) - - q0[c1] *= np.sign(ts[c1,2,1] - ts[c1,1,2]) - q2[c1] *= np.sign(ts[c1,1,0] + ts[c1,0,1]) - q3[c1] *= np.sign(ts[c1,0,2] + ts[c1,2,0]) - - q0[c2] *= np.sign(ts[c2,0,2] - ts[c2,2,0]) - q1[c2] *= np.sign(ts[c2,1,0] + ts[c2,0,1]) - q3[c2] *= np.sign(ts[c2,2,1] + ts[c2,1,2]) - - q0[c3] *= np.sign(ts[c3,1,0] - ts[c3,0,1]) - q1[c3] *= np.sign(ts[c3,2,0] + ts[c3,0,2]) - q2[c3] *= np.sign(ts[c3,2,1] + ts[c3,1,2]) - + + q1[c0] *= np.sign(ts[c0, 2, 1] - ts[c0, 1, 2]) + q2[c0] *= np.sign(ts[c0, 0, 2] - ts[c0, 2, 0]) + q3[c0] *= np.sign(ts[c0, 1, 0] - ts[c0, 0, 1]) + + q0[c1] *= np.sign(ts[c1, 2, 1] - ts[c1, 1, 2]) + q2[c1] *= np.sign(ts[c1, 1, 0] + ts[c1, 0, 1]) + q3[c1] *= np.sign(ts[c1, 0, 2] + ts[c1, 2, 0]) + + q0[c2] *= np.sign(ts[c2, 0, 2] - ts[c2, 2, 0]) + q1[c2] *= np.sign(ts[c2, 1, 0] + ts[c2, 0, 1]) + q3[c2] *= np.sign(ts[c2, 2, 1] + ts[c2, 1, 2]) + + q0[c3] *= np.sign(ts[c3, 1, 0] - ts[c3, 0, 1]) + q1[c3] *= np.sign(ts[c3, 2, 0] + ts[c3, 0, 2]) + q2[c3] *= np.sign(ts[c3, 2, 1] + ts[c3, 1, 2]) + qs = np.empty(ts.shape[:-2] + (4,)) - qs[...,0] = q0 - qs[...,1] = q1 - qs[...,2] = q2 - qs[...,3] = q3 - + qs[..., 0] = q0 + qs[..., 1] = q1 + qs[..., 2] = q2 + qs[..., 3] = q3 + return cls(qs) - - - diff --git a/utils/__init__.py b/utils/__init__.py index abd7b26..55e8919 100644 --- a/utils/__init__.py +++ b/utils/__init__.py @@ -1,6 +1,5 @@ import sys import os + BASEPATH = os.path.dirname(__file__) sys.path.insert(0, BASEPATH) - - diff --git a/utils/animation_2d_data.py b/utils/animation_2d_data.py index d69f34c..6c6e64f 100644 --- a/utils/animation_2d_data.py +++ b/utils/animation_2d_data.py @@ -1,9 +1,10 @@ import sys import os from os.path import join as pjoin + BASEPATH = os.path.dirname(__file__) sys.path.insert(0, BASEPATH) -sys.path.insert(0, pjoin(BASEPATH, '..')) +sys.path.insert(0, pjoin(BASEPATH, "..")) import numpy as np import json @@ -48,14 +49,24 @@ def from_openpose_json(cls, json_dir, scale=0.07, smooth=True): motion = [] joint_map = { 0: 8, - 1: 12, 2: 13, 3: 14, 4: 19, - 5: 9, 6: 10, 7: 11, 8: 22, + 1: 12, + 2: 13, + 3: 14, + 4: 19, + 5: 9, + 6: 10, + 7: 11, + 8: 22, # 9 is somewhere between 0 & 10 10: 1, # 11 is somewhere between 10 and 12 12: 0, - 13: 5, 14: 6, 15: 7, # 16 is a little bit further - 17: 2, 18: 3, 19: 4, # 20 is a little bit further + 13: 5, + 14: 6, + 15: 7, # 16 is a little bit further + 17: 2, + 18: 3, + 19: 4, # 20 is a little bit further } num_joints = 21 @@ -64,23 +75,35 @@ def from_openpose_json(cls, json_dir, scale=0.07, smooth=True): for path in json_files: with open(path) as f: joint_dict = json.load(f) - if len(joint_dict['people']) == 0: + if len(joint_dict["people"]) == 0: if start: raw_joint = motion[-1] motion.append(raw_joint) else: continue start = True - body_joint = np.array(joint_dict['people'][0]['pose_keypoints_2d']).reshape((-1, 3))[:, :2] - lhand_joint = np.array(joint_dict['people'][0]['hand_left_keypoints_2d']).reshape((-1, 3))[:, :2] - rhand_joint = np.array(joint_dict['people'][0]['hand_right_keypoints_2d']).reshape((-1, 3))[:, :2] - raw_joint = np.concatenate([body_joint, lhand_joint, rhand_joint], axis=-2) + body_joint = np.array( + joint_dict["people"][0]["pose_keypoints_2d"] + ).reshape((-1, 3))[:, :2] + lhand_joint = np.array( + joint_dict["people"][0]["hand_left_keypoints_2d"] + ).reshape((-1, 3))[:, :2] + rhand_joint = np.array( + joint_dict["people"][0]["hand_right_keypoints_2d"] + ).reshape((-1, 3))[:, :2] + raw_joint = np.concatenate( + [body_joint, lhand_joint, rhand_joint], axis=-2 + ) if len(motion) > 0: - raw_joint[np.where(raw_joint == 0)] = motion[-1][np.where(raw_joint == 0)] + raw_joint[np.where(raw_joint == 0)] = motion[-1][ + np.where(raw_joint == 0) + ] motion.append(raw_joint) for i in range(len(motion) - 1, 0, -1): - motion[i - 1][np.where(motion[i - 1] == 0)] = motion[i][np.where(motion[i - 1] == 0)] + motion[i - 1][np.where(motion[i - 1] == 0)] = motion[i][ + np.where(motion[i - 1] == 0) + ] motion = np.stack(motion, axis=0) # motion: [T, J, 2] @@ -119,11 +142,13 @@ def test(): bla = {} for num in [27, 32, 95]: - anim2d = AnimationData2D.from_openpose_json(f'../../data/treadmill/json_inputs/{num}') + anim2d = AnimationData2D.from_openpose_json( + f"../../data/treadmill/json_inputs/{num}" + ) bla[str(num)] = {"motion": anim2d.get_projection(), "foot_contact": None} visualize(bla) -if __name__ == '__main__': +if __name__ == "__main__": test() diff --git a/utils/animation_data.py b/utils/animation_data.py index d7a51af..89bc78d 100644 --- a/utils/animation_data.py +++ b/utils/animation_data.py @@ -4,7 +4,7 @@ BASEPATH = os.path.dirname(__file__) sys.path.insert(0, BASEPATH) -sys.path.insert(0, pjoin(BASEPATH, '..')) +sys.path.insert(0, pjoin(BASEPATH, "..")) import argparse import numpy as np @@ -28,11 +28,11 @@ def forward_rotations(skel, rotations, rtpos=None, trim=True): for i, pi in enumerate(skel.topology): if pi == -1: continue - glb[..., i, :] = np.matmul(transforms[..., pi, :, :], - skel.offset[i]) + glb[..., i, :] = np.matmul(transforms[..., pi, :, :], skel.offset[i]) glb[..., i, :] += glb[..., pi, :] - transforms[..., i, :, :] = np.matmul(transforms[..., pi, :, :], - transforms[..., i, :, :]) + transforms[..., i, :, :] = np.matmul( + transforms[..., pi, :, :], transforms[..., i, :, :] + ) if trim: glb = glb[..., skel.chosen_joints, :] return glb @@ -47,23 +47,11 @@ def rotate_coordinates(local3d, angles): cx, cy, cz = np.cos(angles) sx, sy, sz = np.sin(angles) - mat33_x = np.array([ - [1, 0, 0], - [0, cx, sx], - [0, -sx, cx] - ], dtype='float') + mat33_x = np.array([[1, 0, 0], [0, cx, sx], [0, -sx, cx]], dtype="float") - mat33_y = np.array([ - [cy, 0, sy], - [0, 1, 0], - [-sy, 0, cy] - ], dtype='float') + mat33_y = np.array([[cy, 0, sy], [0, 1, 0], [-sy, 0, cy]], dtype="float") - mat33_z = np.array([ - [cz, sz, 0], - [-sz, cz, 0], - [0, 0, 1] - ], dtype='float') + mat33_z = np.array([[cz, sz, 0], [-sz, cz, 0], [0, 0, 1]], dtype="float") local3d = local3d @ mat33_x @ mat33_y @ mat33_z return local3d @@ -97,15 +85,17 @@ def motion_projection(motion, local_x, view_angle=None): output: motion_proj [(B,) J * 2, T] """ - local = get_local3d(local_x, view_angle) # [(B,) 3, 3] + local = get_local3d(local_x, view_angle) # [(B,) 3, 3] T = motion.shape[-1] # proj on xy-plane # motion_proj = (local[[0, 1], :] @ motion) this used to be [2, 3] @ [J, 3, T] # but T doesn't matter here ... what we care is the "3", using [T, J, 3, 1] would also be OK - motion = motion[..., np.newaxis] # [(B,) T, J, 3, 1] - motion_proj = local[..., np.newaxis, np.newaxis, [0, 1], :] @ motion # [(B,), 1, 1, 2, 3] @ [(B,), T, J, 3, 1] => [(B,), T, J, 2, 1] - motion_proj = motion_proj.reshape(motion_proj.shape[:-3] + (-1, )) # [(B,) T, -1] + motion = motion[..., np.newaxis] # [(B,) T, J, 3, 1] + motion_proj = ( + local[..., np.newaxis, np.newaxis, [0, 1], :] @ motion + ) # [(B,), 1, 1, 2, 3] @ [(B,), T, J, 3, 1] => [(B,), T, J, 2, 1] + motion_proj = motion_proj.reshape(motion_proj.shape[:-3] + (-1,)) # [(B,) T, -1] motion_proj = motion_proj.swapaxes(-1, -2) # [(B,) J * 2, T] return motion_proj @@ -120,7 +110,9 @@ def foot_contact_from_positions(positions, fid_l=(3, 4), fid_r=(7, 8)): velfactor = np.array([0.05, 0.05]) feet_contact = [] for fid_index in [fid_l, fid_r]: - foot_vel = (positions[1:, fid_index] - positions[:-1, fid_index]) ** 2 # [T - 1, 2, 3] + foot_vel = ( + positions[1:, fid_index] - positions[:-1, fid_index] + ) ** 2 # [T - 1, 2, 3] foot_vel = np.sum(foot_vel, axis=-1) # [T - 1, 2] foot_contact = (foot_vel < velfactor).astype(np.float) feet_contact.append(foot_contact) @@ -149,9 +141,9 @@ def phase_from_ft(foot_contact, is_debug=False): """ for j in range(2): for i in range(1, total_length): - ft_start[i, j] = (ft[i - 1, j] == 0 and ft[i, j] == 1) + ft_start[i, j] = ft[i - 1, j] == 0 and ft[i, j] == 1 if is_debug: - print('ft_start,', ft_start) + print("ft_start,", ft_start) last, beg_i = -1, -1 starts = [] @@ -167,7 +159,10 @@ def phase_from_ft(foot_contact, is_debug=False): avg_circle = 0 if num_circles == 0 else circle_length * 1.0 / num_circles if is_debug: - print("%d circles, total length = %d, avg length = %.3lf" % (num_circles, circle_length, avg_circle)) + print( + "%d circles, total length = %d, avg length = %.3lf" + % (num_circles, circle_length, avg_circle) + ) if len(starts) == 0: # phase never changed return phases @@ -191,7 +186,7 @@ def phase_from_ft(foot_contact, is_debug=False): phases *= np.pi if is_debug: - print('phases:', phases) + print("phases:", phases) return phases @@ -201,8 +196,12 @@ def across_from_glb(positions, hips=(2, 6), sdrs=(14, 18)): hips, sdrs: left/right hip joints, left/right shoulder joints output: local x-axis for each frame [T, 3] """ - across = positions[..., hips[0], :] - positions[..., hips[1], :] + \ - positions[..., sdrs[0], :] - positions[..., sdrs[1], :] # [T, 3] + across = ( + positions[..., hips[0], :] + - positions[..., hips[1], :] + + positions[..., sdrs[0], :] + - positions[..., sdrs[1], :] + ) # [T, 3] across = across / np.sqrt((across ** 2).sum(axis=-1))[..., np.newaxis] return across @@ -217,12 +216,18 @@ def y_rotation_from_positions(positions, hips=(2, 6), sdrs=(14, 18)): across = across_from_glb(positions, hips=hips, sdrs=sdrs) direction_filterwidth = 20 forward = np.cross(across, np.array([[0, 1, 0]])) - forward = filters.gaussian_filter1d(forward, direction_filterwidth, axis=0, mode='nearest') + forward = filters.gaussian_filter1d( + forward, direction_filterwidth, axis=0, mode="nearest" + ) forward = forward / np.sqrt((forward ** 2).sum(axis=-1))[..., np.newaxis] - target = np.tile(np.array([0, 0, 1]), forward.shape[:-1] + (1, )) - quaters = Quaternions.between(forward, target)[..., np.newaxis, :] # [T, 4] -> [T, 1, 4] - pivots = Pivots.from_quaternions(-quaters).ps # from "target"[0, 0, 1] to current facing direction "forward" + target = np.tile(np.array([0, 0, 1]), forward.shape[:-1] + (1,)) + quaters = Quaternions.between(forward, target)[ + ..., np.newaxis, : + ] # [T, 4] -> [T, 1, 4] + pivots = Pivots.from_quaternions( + -quaters + ).ps # from "target"[0, 0, 1] to current facing direction "forward" return quaters, pivots @@ -232,25 +237,37 @@ class AnimationData: Skeleton [T, Jo * 4 + 4 global params + 4 foot_contact] """ - def __init__(self, full, skel=None, frametime=1/30): + + def __init__(self, full, skel=None, frametime=1 / 30): if skel is None: skel = Skel() self.skel = skel self.frametime = frametime self.len = len(full) self.rotations = full[:, :-8].reshape(self.len, -1, 4) # [T, Jo, 4] - assert self.rotations.shape[1] == len(self.skel.topology), "Rotations do not match the skeleton." + assert self.rotations.shape[1] == len( + self.skel.topology + ), "Rotations do not match the skeleton." self.rotations /= np.sqrt(np.sum(self.rotations ** 2, axis=-1))[..., np.newaxis] self.rt_pos = full[:, -8:-5] # [T, 3] self.rt_rot = full[:, -5:-4] # [T, 1] self.foot_contact = full[:, -4:] # [T, 4] - self.full = np.concatenate([self.rotations.reshape(self.len, -1), self.rt_pos, self.rt_rot, self.foot_contact], axis=-1) + self.full = np.concatenate( + [ + self.rotations.reshape(self.len, -1), + self.rt_pos, + self.rt_rot, + self.foot_contact, + ], + axis=-1, + ) self.phases = None # [T, 1] self.local_x = None # [3] - self.positions_for_proj = None # [T, (J - 1) + 1, 3], trimmed and not forward facing + self.positions_for_proj = ( + None # [T, (J - 1) + 1, 3], trimmed and not forward facing + ) self.global_positions = None - def get_full(self): return self.full @@ -264,7 +281,9 @@ def get_original_rotations(self, rt_rot=None): rt_rotations = Quaternions(self.rotations[:, :1]) # [T, 1, 4] rt_rotations = np.array(yaxis_rotations * rt_rotations) rt_rotations /= np.sqrt((rt_rotations ** 2).sum(axis=-1))[..., np.newaxis] - return np.concatenate((rt_rotations, self.rotations[:, 1:]), axis=1) # [T, J, 4] + return np.concatenate( + (rt_rotations, self.rotations[:, 1:]), axis=1 + ) # [T, J, 4] def get_foot_contact(self, transpose=False): if transpose: @@ -286,18 +305,28 @@ def get_local_x(self): def get_content_input(self): rotations = self.rotations.reshape(self.len, -1) # [T, Jo * 4] - return np.concatenate((rotations, self.rt_pos, self.rt_rot), axis=-1).transpose(1, 0) # [Jo * 4 + 3 + 1, T] + return np.concatenate((rotations, self.rt_pos, self.rt_rot), axis=-1).transpose( + 1, 0 + ) # [Jo * 4 + 3 + 1, T] def get_style3d_input(self): - pos3d = forward_rotations(self.skel, self.rotations, trim=True)[:, 1:] # [T, J - 1, 3] + pos3d = forward_rotations(self.skel, self.rotations, trim=True)[ + :, 1: + ] # [T, J - 1, 3] pos3d = pos3d.reshape((len(pos3d), -1)) # [T, (J - 1) * 3] - return np.concatenate((pos3d, self.rt_pos, self.rt_rot), axis=-1).transpose(1, 0) # [(J - 1) * 3 + 3 + 1, T] + return np.concatenate((pos3d, self.rt_pos, self.rt_rot), axis=-1).transpose( + 1, 0 + ) # [(J - 1) * 3 + 3 + 1, T] def get_projections(self, view_angles, scales=None): if self.positions_for_proj is None: rotations = self.get_original_rotations() - positions = forward_rotations(self.skel, rotations, trim=True)[:, 1:] # [T, J - 1, 3] - positions = np.concatenate((positions, self.rt_pos[:, np.newaxis, :]), axis=1) # [T, J, 3] + positions = forward_rotations(self.skel, rotations, trim=True)[ + :, 1: + ] # [T, J - 1, 3] + positions = np.concatenate( + (positions, self.rt_pos[:, np.newaxis, :]), axis=1 + ) # [T, J, 3] self.positions_for_proj = positions.copy() else: positions = self.positions_for_proj.copy() @@ -305,21 +334,31 @@ def get_projections(self, view_angles, scales=None): if scales is None: scales = np.ones((len(view_angles))) for angle, scale in zip(view_angles, scales): - projections.append(motion_projection(positions, self.get_local_x(), angle) * scale) + projections.append( + motion_projection(positions, self.get_local_x(), angle) * scale + ) projections = np.stack(projections, axis=-3) # [V, J * 2, T] return projections def get_global_positions(self, trim=True): # for visualization if not trim: - return forward_rotations(self.skel, self.get_original_rotations(), rtpos=self.rt_pos, trim=False) + return forward_rotations( + self.skel, self.get_original_rotations(), rtpos=self.rt_pos, trim=False + ) if self.global_positions is None: rotations = self.get_original_rotations() - positions = forward_rotations(self.skel, rotations, rtpos=self.rt_pos, trim=True) + positions = forward_rotations( + self.skel, rotations, rtpos=self.rt_pos, trim=True + ) self.global_positions = positions return self.global_positions def get_velocity_factor(self): - positions = forward_rotations(self.skel, self.get_original_rotations(), trim=True)[:, 1:] # [T, J - 1, 3] + positions = forward_rotations( + self.skel, self.get_original_rotations(), trim=True + )[ + :, 1: + ] # [T, J - 1, 3] velocity = positions[1:] - positions[:-1] # [T - 1, J - 1, 3] velocity = np.sqrt(np.sum(velocity ** 2, axis=-1)) # [T - 1, J - 1] max_velocity = np.max(velocity, axis=-1) # [T - 1] @@ -329,10 +368,14 @@ def get_velocity_factor(self): def get_BVH(self, forward=True): rt_pos = self.rt_pos # [T, 3] rt_rot = self.rt_rot # [T, 1] - if forward: # choose a direction in [z+, x+, z-, x-], which is closest to "forward", as the new z+ + if ( + forward + ): # choose a direction in [z+, x+, z-, x-], which is closest to "forward", as the new z+ directions = np.array(range(4)) * np.pi * 0.5 # [0, 1, 2, 3] * 0.5pi - diff = rt_rot[np.newaxis, :] - directions[:, np.newaxis, np.newaxis] # [1, T, 1] - [4, 1, 1] + diff = ( + rt_rot[np.newaxis, :] - directions[:, np.newaxis, np.newaxis] + ) # [1, T, 1] - [4, 1, 1] diff = np.minimum(np.abs(diff), 2.0 * np.pi - np.abs(diff)) diff = np.sum(diff, axis=(-1, -2)) # [4, T, 1] -> [4] @@ -361,7 +404,9 @@ def from_network_output(cls, input): return cls(input) @classmethod - def from_rotations_and_root_positions(cls, rotations, root_positions, skel=None, frametime=1/30): + def from_rotations_and_root_positions( + cls, rotations, root_positions, skel=None, frametime=1 / 30 + ): """ rotations: [T, J, 4] root_positions: [T, 3] @@ -371,15 +416,27 @@ def from_rotations_and_root_positions(cls, rotations, root_positions, skel=None, rotations /= np.sqrt(np.sum(rotations ** 2, axis=-1))[..., np.newaxis] global_positions = forward_rotations(skel, rotations, root_positions, trim=True) - foot_contact = foot_contact_from_positions(global_positions, fid_l=skel.fid_l, fid_r=skel.fid_r) - quaters, pivots = y_rotation_from_positions(global_positions, hips=skel.hips, sdrs=skel.sdrs) + foot_contact = foot_contact_from_positions( + global_positions, fid_l=skel.fid_l, fid_r=skel.fid_r + ) + quaters, pivots = y_rotation_from_positions( + global_positions, hips=skel.hips, sdrs=skel.sdrs + ) root_rotations = Quaternions(rotations[:, 0:1, :].copy()) # [T, 1, 4] root_rotations = quaters * root_rotations # facing [0, 0, 1] root_rotations = np.array(root_rotations).reshape((-1, 1, 4)) # [T, 1, 4] rotations[:, 0:1, :] = root_rotations - full = np.concatenate([rotations.reshape((len(rotations), -1)), root_positions, pivots, foot_contact], axis=-1) + full = np.concatenate( + [ + rotations.reshape((len(rotations), -1)), + root_positions, + pivots, + foot_contact, + ], + axis=-1, + ) return cls(full, skel, frametime) @classmethod @@ -391,19 +448,20 @@ def from_BVH(cls, filename, downsample=4, skel=None, trim_scale=None): anim = anim[:length] rotations = np.array(anim.rotations) # [T, J, 4] root_positions = anim.positions[:, 0, :] - return cls.from_rotations_and_root_positions(rotations, root_positions, skel=skel, frametime=frametime * downsample) + return cls.from_rotations_and_root_positions( + rotations, root_positions, skel=skel, frametime=frametime * downsample + ) def parse_args(): parser = argparse.ArgumentParser("test") - parser.add_argument('--bvh_in', type=str, default=None) - parser.add_argument('--dataset', type=str, default=None) + parser.add_argument("--bvh_in", type=str, default=None) + parser.add_argument("--dataset", type=str, default=None) return parser.parse_args() def test_all(args): - def mse(a, b): return np.sum((a - b) ** 2) @@ -412,7 +470,9 @@ def test_phase_from_ft(): pace[::8] = 1 left = pace[:-4] right = pace[4:] - phase_from_ft(np.concatenate([left, left, right, right], axis=-1), is_debug=True) + phase_from_ft( + np.concatenate([left, left, right, right], axis=-1), is_debug=True + ) def BVH_and_back(filename): anim, names, frametime = BVH.load(filename) @@ -421,23 +481,27 @@ def BVH_and_back(filename): root_positions = anim.positions[:, 0, :] anim_a = AnimationData.from_BVH(filename) - rotations = rotations / np.sqrt(np.sum(rotations ** 2, axis=-1))[..., np.newaxis] - print(f'rotations: {mse(anim_a.get_original_rotations(), rotations)}') - print(f'root_positions: {mse(anim_a.get_root_positions(), root_positions)}') + rotations = ( + rotations / np.sqrt(np.sum(rotations ** 2, axis=-1))[..., np.newaxis] + ) + print(f"rotations: {mse(anim_a.get_original_rotations(), rotations)}") + print(f"root_positions: {mse(anim_a.get_root_positions(), root_positions)}") content_input = anim_a.get_content_input() style3d_input = anim_a.get_style3d_input() view_points = () for i in range(7): - view_points += ((0, -np.pi / 2 + i * np.pi / 6, 0), ) + view_points += ((0, -np.pi / 2 + i * np.pi / 6, 0),) view_points = () scales = () for i in range(4): - view_points += ((0, -np.pi / 2 + float(np.random.rand(1)) * np.pi, 0), ) - scales += (float(np.random.rand(1)) * 0.4 + 0.8, ) + view_points += ((0, -np.pi / 2 + float(np.random.rand(1)) * np.pi, 0),) + scales += (float(np.random.rand(1)) * 0.4 + 0.8,) style2d_input = anim_a.get_projections(view_points, scales) - print(f'content {content_input.shape}, style3d {style3d_input.shape}, style2d {style2d_input.shape}') + print( + f"content {content_input.shape}, style3d {style3d_input.shape}, style2d {style2d_input.shape}" + ) foot_contact = anim_a.get_foot_contact() T = content_input.shape[-1] @@ -446,7 +510,9 @@ def BVH_and_back(filename): inplace = anim_a.positions_for_proj[:, :-1, :] inplace = np.concatenate((np.zeros((T, 1, 3)), inplace), axis=1) original = anim_a.get_global_positions() - print(f'inplace no rot {inplace_no_rot.shape}, inplace {inplace.shape}, original {original.shape}') + print( + f"inplace no rot {inplace_no_rot.shape}, inplace {inplace.shape}, original {original.shape}" + ) """ visualize({ @@ -457,12 +523,18 @@ def BVH_and_back(filename): """ motion_proj = {} - for (view_point, scale, proj) in zip(view_points, scales, style2d_input): # [V, J * 2, T] + for (view_point, scale, proj) in zip( + view_points, scales, style2d_input + ): # [V, J * 2, T] proj = proj.copy().transpose(1, 0).reshape(T, -1, 2) # [T, J, 2] proj = np.concatenate([proj[:, -1:], proj[:, :-1]], axis=1) - ori_proj = np.concatenate([proj[:, :1], proj[:, 1:] + proj[:, :1].copy()], axis=1) + ori_proj = np.concatenate( + [proj[:, :1], proj[:, 1:] + proj[:, :1].copy()], axis=1 + ) proj[:, :1] = 0 - motion_proj[f'angle: {(view_point[1] / np.pi * 180):3f} scale: {scale:3f}'] = {"motion": ori_proj, "foot_contact": foot_contact} + motion_proj[ + f"angle: {(view_point[1] / np.pi * 180):3f} scale: {scale:3f}" + ] = {"motion": ori_proj, "foot_contact": foot_contact} """ visualize({ "inplace_proj": {"motion": proj, "foot_contact": foot_contact}, @@ -478,7 +550,9 @@ def check_velocity(dataset): motions, labels, metas = dataset["motion"], dataset["style"], dataset["meta"] style_names = list(set(metas["style"])) content_names = list(set(metas["content"])) - info = {content: {style: [] for style in style_names} for content in content_names} + info = { + content: {style: [] for style in style_names} for content in content_names + } for i, motion in enumerate(motions): anim = AnimationData(motion, skel=skel) vel = anim.get_velocity_factor() @@ -492,11 +566,11 @@ def check_velocity(dataset): info[content]["all"] = np.mean(all) with open("probe_velocity.csv", "w") as f: - columns = ['all'] + style_names - f.write(',' + ','.join(columns) + '\n') + columns = ["all"] + style_names + f.write("," + ",".join(columns) + "\n") for content in info: - values = [f'{info[content][key]}' for key in columns] - f.write(','.join([content] + values) + '\n') + values = [f"{info[content][key]}" for key in columns] + f.write(",".join([content] + values) + "\n") dataset = np.load(args.dataset, allow_pickle=True)["trainfull"].item() check_velocity(dataset) @@ -504,7 +578,6 @@ def check_velocity(dataset): # BVH_and_back(args.bvh_in) -if __name__ == '__main__': +if __name__ == "__main__": args = parse_args() test_all(args) - diff --git a/utils/load_skeleton.py b/utils/load_skeleton.py index b9af7fd..be3feaf 100644 --- a/utils/load_skeleton.py +++ b/utils/load_skeleton.py @@ -3,6 +3,7 @@ import os import sys from os.path import join as pjoin + BASEPATH = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, BASEPATH) import BVH as BVH @@ -10,26 +11,30 @@ class Skel: - def __init__(self, filename=pjoin(BASEPATH, "..", "style_transfer", "global_info", "skeleton_CMU.yml")): + def __init__( + self, + filename=pjoin( + BASEPATH, "..", "style_transfer", "global_info", "skeleton_CMU.yml" + ), + ): f = open(filename, "r") skel = yaml.load(f, Loader=yaml.Loader) - self.bvh_name = pjoin(os.path.dirname(filename), skel['BVH']) + self.bvh_name = pjoin(os.path.dirname(filename), skel["BVH"]) self.rest_bvh = BVH.load(self.bvh_name) - self.offset = np.array(skel['offsets']) - self.topology = np.array(skel['parents']) - self.chosen_joints = np.array(skel['chosen_joints']) - self.chosen_parents = np.array(skel['chosen_parents']) - self.fid_l, self.fid_r = skel['left_foot'], skel['right_foot'] - self.hips, self.sdrs = skel['hips'], skel['shoulders'] - self.head = skel['head'] - self.visualization = skel['visualization'] + self.offset = np.array(skel["offsets"]) + self.topology = np.array(skel["parents"]) + self.chosen_joints = np.array(skel["chosen_joints"]) + self.chosen_parents = np.array(skel["chosen_parents"]) + self.fid_l, self.fid_r = skel["left_foot"], skel["right_foot"] + self.hips, self.sdrs = skel["hips"], skel["shoulders"] + self.head = skel["head"] + self.visualization = skel["visualization"] -if __name__ == '__main__': +if __name__ == "__main__": skel = Skel() print(skel.topology) print(skel.offset) print(skel.rest_bvh[0].offsets) print(skel.chosen_joints) print(skel.chosen_parents) -