def set_new_root(self, new_root): euler = torch.tensor(self.anim.rotations[:, 0, :], dtype=torch.float) 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.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') new_rot1 = 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 self.anim.rotations[:, new_root, :] = new_rot1 new_seq = [] vis = [0] * self.anim.rotations.shape[1] new_idx = [-1] * len(vis) new_parent = [0] * len(vis) def relabel(x): nonlocal new_seq, vis, new_idx, new_parent new_idx[x] = len(new_seq) 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): relabel(y) new_parent[new_idx[y]] = new_idx[x] relabel(new_root) self.anim.rotations = self.anim.rotations[:, new_seq, :] self.anim.offsets = self.anim.offsets[new_seq] names = self._names.copy() for i, j in enumerate(new_seq): self._names[i] = names[j] self.anim.parents = np.array(new_parent, dtype=np.int)
def get_root_rotation(self): """ Get the rotation of the root joint :return: A numpy Array in shape [frame, 1, 4] """ # rotation in shape[frame, 1, 4] rotation = self.anim.rotations[:, self.corps, :] rotation = rotation[:, 0:1, :] # convert to quaternion format rotation = Quaternions.from_euler(np.radians(rotation)).qs return rotation
def rotate(self, theta, 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_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] self.anim.positions[:, 0, :] = position
def get_rotation(self): """ Get the rotation of each joint's parent except the root joint :return: A numpy Array in shape [frame, simple_joint_num - 1, 4]] """ # rotation in shape[frame, simple_joint_num, 3] rotations = self.anim.rotations[:, self.corps, :] # transform euler degree into radians then into quaternion # in shape [frame, simple_joint_num, 4] rotations = Quaternions.from_euler(np.radians(rotations)).qs # 除去root的剩下所有joint信息都存储在edges中 # 因为.bvh中每个joint的rotation指的是施加在child joint上的旋转 # 而此时简化的骨架的旋转是完整骨架中每根骨骼的parent的旋转,所以要把这些旋转信息替换成 # 简化骨骼中对应的parent的旋转 # 又因为使用了edges的信息,所以一定程度上可以看作这些edges自身旋转了多少 index = [] for e in self.edges: index.append(e[0]) # now rotation is in shape[frame, simple_joint_num - 1, 4] rotations = rotations[:, index, :] return rotations
def load_bvh(file_name, channel_map=None, start=None, end=None, order=None, world=False): ''' Reads a BVH file and constructs an animation Parameters ---------- file_name: str File to be opened channel_map: Dict Mapping between the coordinates x, y, z and the positions X, Y, Z in the bvh file start : int Optional Starting Frame end : int Optional Ending Frame order : str Optional Specifier for joint order. Given as string E.G 'xyz', 'zxy' world : bool If set to true euler angles are applied together in world space rather than local space Returns ------- (animation, joint_names, frame_time) Tuple of loaded animation and joint names ''' if channel_map is None: channel_map = { 'Xrotation': 'x', 'Yrotation': 'y', 'Zrotation': 'z' } f = open(file_name, 'r') i = 0 active = -1 end_site = False names = [] orients = Quaternions.id(0) offsets = np.array([]).reshape((0, 3)) parents = np.array([], dtype=int) prev_data_match = [] for line in f: if 'HIERARCHY' in line: continue if 'MOTION' in line: continue root_match = re.match(r'ROOT (\w+)', line) if root_match: names.append(root_match.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) continue if '{' in line: continue if '}' in line: if end_site: end_site = False else: active = parents[active] continue offset_match = re.match( r'\s*OFFSET\s+([\-\d\.e]+)\s+([\-\d\.e]+)\s+([\-\d\.e]+)', line) if offset_match: if not end_site: offsets[active] = np.array( [list(map(float, offset_match.groups()))]) continue channel_match = re.match(r'\s*CHANNELS\s+(\d+)', line) if channel_match: channels = int(channel_match.group(1)) if order is None: channel_is = 0 if channels == 3 else 3 channel_ie = 3 if channels == 3 else 6 parts = line.split()[2 + channel_is:2 + channel_ie] if any([p not in channel_map for p in parts]): continue order = ''.join([channel_map[p] for p in parts]) continue joint_match = re.match('\s*JOINT\s+(\w+)', line) if joint_match: names.append(joint_match.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) continue if 'End Site' in line: end_site = True continue frame_match = re.match('\s*Frames:\s+(\d+)', line) if frame_match: if start and end: frame_num = (end - start) - 1 else: frame_num = int(frame_match.group(1)) joint_num = len(parents) positions = offsets[np.newaxis].repeat(frame_num, axis=0) rotations = np.zeros((frame_num, len(orients), 3)) continue frame_match = re.match('\s*Frame Time:\s+([\d\.]+)', line) if frame_match: frame_time = float(frame_match.group(1)) continue if (start and end) and (i < start or i >= end - 1): i += 1 continue if '#IND' in line: new_data_match = line.strip().split(' ') for i in range(len(new_data_match)): if '#IND' in new_data_match[i]: new_data_match[i] = prev_data_match[i] line = ' '.join(new_data_match) data_match = line.strip().split(' ') prev_data_match = data_match if data_match: data_block = np.array(list(map(float, data_match))) 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) elif channels == 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] else: raise Exception('Too many channels! {}'.format(channels)) i += 1 f.close() rotations = qfix( Quaternions.from_euler(np.radians(rotations), order=order, world=world).qs) positions = MocapDataset.forward_kinematics( torch.from_numpy(rotations).cuda().float().unsqueeze(0), torch.from_numpy(positions[:, 0]).cuda().float().unsqueeze(0), parents, torch.from_numpy(offsets).cuda().float()).squeeze().cpu().numpy() orientations, _ = Quaternions(rotations[:, 0]).angle_axis() return names, parents, offsets, positions, rotations
def convert_eular_to_quaternions(self, rotations): rotations_reshape = rotations.view((-1, 3)) q = Quaternions.from_euler(rotations_reshape) return q.qs
def load(filename, start=None, end=None, order=None, world=False): """ Reads a BVH file and constructs an animation Parameters ---------- filename: str File to be opened start : int Optional Starting Frame end : int Optional Ending Frame order : str Optional Specifier for joint order. Given as string E.G 'xyz', 'zxy' world : bool If set to true euler angles are applied together in world space rather than local space Returns ------- (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)) parents = np.array([], dtype=int) for line in f: 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) continue if "{" in line: continue if "}" in line: 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) 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] if any([p not in channelmap for p in parts]): continue order = "".join([channelmap[p] for p in parts]) continue 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) 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 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): 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: 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] 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] 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
def load_from_maya(root, start, end): """ Load Animation Object from Maya Joint Skeleton Parameters ---------- root : PyNode Root Joint of Maya Skeleton start, end : int, int Start and End frame index of Maya Animation Returns ------- animation : Animation Loaded animation from maya names : [str] Joint names from maya """ 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 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) scales = {} for j, name, parent in zip(range(len(names)), names, parents): node = pm.PyNode(name) 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) 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) """ Done """ pm.currentTime(original_time) return Animation(rotations, positions, orients, offsets, parents), names
def generate_motion(self, load_saved_model=True, samples_to_generate=1534, max_steps=300, randomized=True): if load_saved_model: self.load_best_model() self.model.eval() test_loader = self.data_loader['test'] pos, quat, orient, z_mean, z_dev, \ root_speed, affs, spline, labels = self.return_batch([samples_to_generate], test_loader, randomized=randomized) pos = torch.from_numpy(pos).cuda() traj = pos[:, :, 0, [0, 2]].clone() orient = torch.from_numpy(orient).cuda() quat = torch.from_numpy(quat).cuda() z_mean = torch.from_numpy(z_mean).cuda() z_dev = torch.from_numpy(z_dev).cuda() root_speed = torch.from_numpy(root_speed).cuda() affs = torch.from_numpy(affs).cuda() spline = torch.from_numpy(spline).cuda() z_rs = torch.cat((z_dev, root_speed), dim=-1) quat_all = torch.cat((orient[:, self.prefix_length - 1:], quat[:, self.prefix_length - 1:]), dim=-1) labels = np.tile(labels, (1, max_steps + self.prefix_length, 1)) # Begin for transition # traj[:, self.prefix_length - 2] = torch.tensor([-0.208, 4.8]).cuda().float() # traj[:, self.prefix_length - 1] = torch.tensor([-0.204, 5.1]).cuda().float() # final_emo_idx = int(max_steps/2) # labels[:, final_emo_idx:] = np.array([1., 0., 0., 0.]) # labels[:, :final_emo_idx + 1] = np.linspace(labels[:, 0], labels[:, final_emo_idx], # num=final_emo_idx + 1, axis=1) # End for transition labels = torch.from_numpy(labels).cuda() # traj_np = traj_markers.detach().cpu().numpy() # import matplotlib.pyplot as plt # plt.plot(traj_np[6, :, 0], traj_np[6, :, 1]) # plt.show() happy_idx = [25, 295, 390, 667, 1196] sad_idx = [169, 184, 258, 948, 974] angry_idx = [89, 93, 96, 112, 289, 290, 978] neutral_idx = [72, 106, 143, 237, 532, 747, 1177] sample_idx = np.squeeze(np.concatenate((happy_idx, sad_idx, angry_idx, neutral_idx))) ## CHANGE HERE # scene_corners = torch.tensor([[149.862, 50.833], # [149.862, 36.81], # [161.599, 36.81], # [161.599, 50.833]]).cuda().float() # character_heights = torch.tensor([0.95, 0.88, 0.86, 0.90, 0.95, 0.82]).cuda().float() # num_characters_per_side = torch.tensor([2, 3, 2, 3]).cuda().int() # traj_markers, traj_offsets, character_scale =\ # generate_trajectories(scene_corners, z_mean, character_heights, # num_characters_per_side, traj[:, :self.prefix_length]) # num_characters_per_side = torch.tensor([4, 0, 0, 0]).cuda().int() # traj_markers, traj_offsets, character_scale =\ # generate_simple_trajectories(scene_corners, z_mean[:4], z_mean[:4], # num_characters_per_side, traj[sample_idx, :self.prefix_length]) # traj_markers, traj_offsets, character_scale =\ # generate_rvo_trajectories(scene_corners, z_mean[:4], z_mean[:4], # num_characters_per_side, traj[sample_idx, :self.prefix_length]) # traj[sample_idx, :self.prefix_length] += traj_offsets # pos_sampled = pos[sample_idx].clone() # pos_sampled[:, :self.prefix_length, :, [0, 2]] += traj_offsets.unsqueeze(2).repeat(1, 1, self.V, 1) # pos[sample_idx] = pos_sampled # traj_markers = self.generate_linear_trajectory(traj) pos_pred = self.copy_prefix(pos) traj_pred = self.copy_prefix(traj) orient_pred = self.copy_prefix(orient) quat_pred = self.copy_prefix(quat) z_rs_pred = self.copy_prefix(z_rs) affs_pred = self.copy_prefix(affs) spline_pred = self.copy_prefix(spline) labels_pred = self.copy_prefix(labels, prefix_length=max_steps + self.prefix_length) # forward elapsed_time = np.zeros(len(sample_idx)) for counter, s in enumerate(sample_idx): # range(samples_to_generate): start_time = time.time() orient_h_copy = self.orient_h.clone() quat_h_copy = self.quat_h.clone() ## CHANGE HERE num_markers = max_steps + self.prefix_length + 1 # num_markers = traj_markers[s].shape[0] marker_idx = 0 t = -1 with torch.no_grad(): while marker_idx < num_markers: t += 1 if t > max_steps: print('Sample: {}. Did not reach end in {} steps.'.format(s, max_steps), end='') break pos_pred[s] = torch.cat((pos_pred[s], torch.zeros_like(pos_pred[s][:, -1:])), dim=1) traj_pred[s] = torch.cat((traj_pred[s], torch.zeros_like(traj_pred[s][:, -1:])), dim=1) orient_pred[s] = torch.cat((orient_pred[s], torch.zeros_like(orient_pred[s][:, -1:])), dim=1) quat_pred[s] = torch.cat((quat_pred[s], torch.zeros_like(quat_pred[s][:, -1:])), dim=1) z_rs_pred[s] = torch.cat((z_rs_pred[s], torch.zeros_like(z_rs_pred[s][:, -1:])), dim=1) affs_pred[s] = torch.cat((affs_pred[s], torch.zeros_like(affs_pred[s][:, -1:])), dim=1) spline_pred[s] = torch.cat((spline_pred[s], torch.zeros_like(spline_pred[s][:, -1:])), dim=1) orient_pred[s][:, self.prefix_length + t:self.prefix_length + t + 1], \ quat_pred[s][:, self.prefix_length + t:self.prefix_length + t + 1], \ z_rs_pred[s][:, self.prefix_length + t:self.prefix_length + t + 1], \ orient_h_copy, quat_h_copy = \ self.model( orient_pred[s][:, t:self.prefix_length + t], quat_pred[s][:, t:self.prefix_length + t], z_rs_pred[s][:, t:self.prefix_length + t], affs_pred[s][:, t:self.prefix_length + t], spline_pred[s][:, t:self.prefix_length + t], labels_pred[s][:, t:self.prefix_length + t], orient_h=None if t == 0 else orient_h_copy, quat_h=None if t == 0 else quat_h_copy, return_prenorm=False) traj_curr = traj_pred[s][:, self.prefix_length + t - 1:self.prefix_length + t].clone() # root_speed = torch.norm( # pos_pred[s][:, self.prefix_length + t:self.prefix_length + t + 1, 0] - \ # pos_pred[s][:, self.prefix_length + t - 1:self.prefix_length + t, 0], dim=-1) ## CHANGE HERE # traj_next = \ # self.compute_next_traj_point( # traj_curr, # traj_markers[s, marker_idx], # o_z_rs_pred[s][:, self.prefix_length + t:self.prefix_length + t + 1, 2]) try: traj_next = traj[s, self.prefix_length + t] except IndexError: traj_next = \ self.compute_next_traj_point_sans_markers( pos_pred[s][:, self.prefix_length + t - 1:self.prefix_length + t], quat_pred[s][:, self.prefix_length + t:self.prefix_length + t + 1], z_rs_pred[s][:, self.prefix_length + t:self.prefix_length + t + 1, 0], z_rs_pred[s][:, self.prefix_length + t:self.prefix_length + t + 1, 1]) pos_pred[s][:, self.prefix_length + t:self.prefix_length + t + 1], \ affs_pred[s][:, self.prefix_length + t:self.prefix_length + t + 1], \ spline_pred[s][:, self.prefix_length + t:self.prefix_length + t + 1] = \ self.mocap.get_predicted_features( pos_pred[s][:, :self.prefix_length + t], traj_next, z_rs_pred[s][:, self.prefix_length + t:self.prefix_length + t + 1, 0] + z_mean[s:s + 1], orient_pred[s][:, self.prefix_length + t:self.prefix_length + t + 1], quat_pred[s][:, self.prefix_length + t:self.prefix_length + t + 1]) # min_speed_pred = torch.min(torch.cat((lf_speed_pred.unsqueeze(-1), # rf_speed_pred.unsqueeze(-1)), dim=-1), dim=-1)[0] # if min_speed_pred - diff_speeds_mean[s] - diff_speeds_std[s] < 0.: # root_speed_pred = 0. # else: # root_speed_pred = o_z_rs_pred[s][:, self.prefix_length + t:self.prefix_length + t + 1, 2] # ## CHANGE HERE # traj_pred[s][:, self.prefix_length + t:self.prefix_length + t + 1] = \ # self.compute_next_traj_point( # traj_curr, # traj_markers[s, marker_idx], # root_speed_pred) # if torch.norm(traj_next - traj_curr, dim=-1).squeeze() >= \ # torch.norm(traj_markers[s, marker_idx] - traj_curr, dim=-1).squeeze(): # marker_idx += 1 traj_pred[s][:, self.prefix_length + t:self.prefix_length + t + 1] = traj_next marker_idx += 1 pos_pred[s][:, self.prefix_length + t:self.prefix_length + t + 1], \ affs_pred[s][:, self.prefix_length + t:self.prefix_length + t + 1], \ spline_pred[s][:, self.prefix_length + t:self.prefix_length + t + 1] = \ self.mocap.get_predicted_features( pos_pred[s][:, :self.prefix_length + t], pos_pred[s][:, self.prefix_length + t:self.prefix_length + t + 1, 0, [0, 2]], z_rs_pred[s][:, self.prefix_length + t:self.prefix_length + t + 1, 0] + z_mean[s:s + 1], orient_pred[s][:, self.prefix_length + t:self.prefix_length + t + 1], quat_pred[s][:, self.prefix_length + t:self.prefix_length + t + 1]) print('Sample: {}. Steps: {}'.format(s, t), end='\r') print() # shift = torch.zeros((1, scene_corners.shape[1] + 1)).cuda().float() # shift[..., [0, 2]] = scene_corners[0] # pos_pred[s] = (pos_pred[s] - shift) / character_scale + shift # pos_pred_np = pos_pred[s].contiguous().view(pos_pred[s].shape[0], # pos_pred[s].shape[1], -1).permute(0, 2, 1).\ # detach().cpu().numpy() # display_animations(pos_pred_np, self.V, self.C, self.mocap.joint_parents, save=True, # dataset_name=self.dataset, subset_name='epoch_' + str(self.best_loss_epoch), # save_file_names=[str(s).zfill(6)], # overwrite=True) # plt.cla() # fig, (ax1, ax2) = plt.subplots(2, 1) # ax1.plot(root_speeds[s]) # ax1.plot(lf_speeds[s]) # ax1.plot(rf_speeds[s]) # ax1.plot(min_speeds[s] - root_speeds[s]) # ax1.legend(['root', 'left', 'right', 'diff']) # ax2.plot(root_speeds_pred) # ax2.plot(lf_speeds_pred) # ax2.plot(rf_speeds_pred) # ax2.plot(min_speeds_pred - root_speeds_pred) # ax2.legend(['root', 'left', 'right', 'diff']) # plt.show() head_tilt = np.tile(np.array([0., 0., 0.]), (1, quat_pred[s].shape[1], 1)) l_shoulder_slouch = np.tile(np.array([0., 0., 0.]), (1, quat_pred[s].shape[1], 1)) r_shoulder_slouch = np.tile(np.array([0., 0., 0.]), (1, quat_pred[s].shape[1], 1)) head_tilt = Quaternions.from_euler(head_tilt, order='xyz').qs l_shoulder_slouch = Quaternions.from_euler(l_shoulder_slouch, order='xyz').qs r_shoulder_slouch = Quaternions.from_euler(r_shoulder_slouch, order='xyz').qs # Begin for aligning facing direction to trajectory axis_diff, angle_diff = self.get_diff_from_traj(pos_pred, traj_pred, s) angle_thres = 0.3 # angle_thres = torch.max(angle_diff[:, 1:self.prefix_length]) angle_diff[angle_diff <= angle_thres] = 0. angle_diff[:, self.prefix_length] = 0. # End for aligning facing direction to trajectory # pos_copy, quat_copy = self.rotate_gaits(pos_pred, quat_pred, quat_diff, # head_tilt, l_shoulder_slouch, r_shoulder_slouch, s) # pos_pred[s] = pos_copy.clone() # angle_diff_intermediate = self.get_diff_from_traj(pos_pred, traj_pred, s) # if torch.max(angle_diff_intermediate[:, self.prefix_length:]) > np.pi / 2.: # quat_diff = Quaternions.from_angle_axis(-angle_diff.cpu().numpy(), np.array([0, 1, 0])).qs # pos_copy, quat_copy = self.rotate_gaits(pos_pred, quat_pred, quat_diff, # head_tilt, l_shoulder_slouch, r_shoulder_slouch, s) # pos_pred[s] = pos_copy.clone() # axis_diff = torch.zeros_like(axis_diff) # axis_diff[..., 1] = 1. # angle_diff = torch.zeros_like(angle_diff) quat_diff = torch.from_numpy(Quaternions.from_angle_axis( angle_diff.cpu().numpy(), axis_diff.cpu().numpy()).qs).cuda().float() orient_pred[s], quat_pred[s] = self.rotate_gaits(orient_pred[s], quat_pred[s], quat_diff, head_tilt, l_shoulder_slouch, r_shoulder_slouch) if labels_pred[s][:, 0, 0] > 0.5: label_dir = 'happy' elif labels_pred[s][:, 0, 1] > 0.5: label_dir = 'sad' elif labels_pred[s][:, 0, 2] > 0.5: label_dir = 'angry' else: label_dir = 'neutral' ## CHANGE HERE # pos_pred[s] = pos_pred[s][:, self.prefix_length + 5:] # o_z_rs_pred[s] = o_z_rs_pred[s][:, self.prefix_length + 5:] # quat_pred[s] = quat_pred[s][:, self.prefix_length + 5:] traj_pred_np = pos_pred[s][0, :, 0].cpu().numpy() save_file_name = '{:06}_{:.2f}_{:.2f}_{:.2f}_{:.2f}'.format(s, labels_pred[s][0, 0, 0], labels_pred[s][0, 0, 1], labels_pred[s][0, 0, 2], labels_pred[s][0, 0, 3]) animation_pred = { 'joint_names': self.joint_names, 'joint_offsets': torch.from_numpy(self.joint_offsets[1:]).float().unsqueeze(0).repeat( len(pos_pred), 1, 1), 'joint_parents': self.joint_parents, 'positions': pos_pred[s], 'rotations': torch.cat((orient_pred[s], quat_pred[s]), dim=-1) } self.mocap.save_as_bvh(animation_pred, dataset_name=self.dataset, # subset_name='epoch_' + str(self.best_loss_epoch), # save_file_names=[str(s).zfill(6)]) subset_name=os.path.join('no_aff_epoch_' + str(self.best_loss_epoch), str(counter).zfill(2) + '_' + label_dir), save_file_names=['root']) end_time = time.time() elapsed_time[counter] = end_time - start_time print('Elapsed Time: {}'.format(elapsed_time[counter])) # display_animations(pos_pred_np, self.V, self.C, self.mocap.joint_parents, save=True, # dataset_name=self.dataset, # # subset_name='epoch_' + str(self.best_loss_epoch), # # save_file_names=[str(s).zfill(6)], # subset_name=os.path.join('epoch_' + str(self.best_loss_epoch), label_dir), # save_file_names=[save_file_name], # overwrite=True) print('Mean Elapsed Time: {}'.format(np.mean(elapsed_time)))