| import re |
| import numpy as np |
| from common.quaternion import * |
| from visualization.Animation import Animation |
|
|
| channelmap = { |
| 'Xrotation': 'x', |
| 'Yrotation': 'y', |
| 'Zrotation': 'z' |
| } |
|
|
| channelmap_inv = { |
| 'x': 'Xrotation', |
| 'y': 'Yrotation', |
| 'z': 'Zrotation', |
| } |
|
|
| ordermap = { |
| 'x': 0, |
| 'y': 1, |
| 'z': 2, |
| } |
|
|
| def load(filename, start=None, end=None, world=False, need_quater=True): |
| """ |
| 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 = Quaterions.id(0) |
| offsets = np.array([]).reshape((0, 3)) |
| parents = np.array([], dtype=int) |
| orders = [] |
|
|
| 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 = np.append(orients, 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)) |
|
|
| 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]) |
| orders.append(order) |
| 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 = np.append(orients, 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() |
|
|
| all_rotations = [] |
| canonical_order = 'xyz' |
| for i, order in enumerate(orders): |
| rot = rotations[:, i:i + 1] |
| if need_quater: |
| quat = euler_to_quat_np(np.radians(rot), order=order, world=world) |
| all_rotations.append(quat) |
| continue |
| elif order != canonical_order: |
| quat = euler_to_quat_np(np.radians(rot), order=order, world=world) |
| rot = np.degrees(qeuler_np(quat, order=canonical_order)) |
| all_rotations.append(rot) |
| rotations = np.concatenate(all_rotations, axis=1) |
|
|
| return Animation(rotations, positions, orients, offsets, parents, names, frametime) |
|
|
| def write_bvh(parent, offset, rotation, rot_position, names, frametime, order, path, endsite=None): |
| file = open(path, 'w') |
| frame = rotation.shape[0] |
| assert rotation.shape[-1] == 3 |
| joint_num = rotation.shape[1] |
| order = order.upper() |
|
|
| file_string = 'HIERARCHY\n' |
|
|
| seq = [] |
|
|
| def write_static(idx, prefix): |
| nonlocal parent, offset, rotation, names, order, endsite, file_string, seq |
| seq.append(idx) |
| if idx == 0: |
| 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' |
|
|
| has_child = False |
| for y in range(idx + 1, rotation.shape[1]): |
| if parent[y] == idx: |
| has_child = True |
| 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 + '}\n' |
|
|
| write_static(0, '') |
|
|
| file_string += 'MOTION\n' + 'Frames: {}\n'.format(frame) + 'Frame Time: %.8f\n' % frametime |
| for i in range(frame): |
| file_string += '%.6f %.6f %.6f ' % (rot_position[i][0], rot_position[i][1], |
| rot_position[i][2]) |
| for j in range(joint_num): |
| idx = seq[j] |
| file_string += '%.6f %.6f %.6f ' % (rotation[i][idx][0], rotation[i][idx][1], rotation[i][idx][2]) |
| file_string += '\n' |
|
|
| file.write(file_string) |
| return file_string |
|
|
| class WriterWrapper: |
| def __init__(self, parents, frametime, offset=None, names=None): |
| self.parents = parents |
| self.offset = offset |
| self.frametime = frametime |
| self.names = names |
|
|
| def write(self, filename, rot, r_pos, order, offset=None, names=None, repr='quat'): |
| """ |
| Write animation to bvh file |
| :param filename: |
| :param rot: Quaternion as (w, x, y, z) |
| :param pos: |
| :param offset: |
| :return: |
| """ |
| if repr not in ['euler', 'quat', 'quaternion', 'cont6d']: |
| raise Exception('Unknown rotation representation') |
| if offset is None: |
| offset = self.offset |
| if not isinstance(offset, torch.Tensor): |
| offset = torch.tensor(offset) |
| n_bone = offset.shape[0] |
|
|
| if repr == 'cont6d': |
| rot = rot.reshape(rot.shape[0], -1, 6) |
| rot = cont6d_to_quat_np(rot) |
| if repr == 'cont6d' or repr == 'quat' or repr == 'quaternion': |
| |
| |
| euler = qeuler_np(rot, order=order) |
| rot = euler |
|
|
| if names is None: |
| if self.names is None: |
| names = ['%02d' % i for i in range(n_bone)] |
| else: |
| names = self.names |
| write_bvh(self.parents, offset, rot, r_pos, names, self.frametime, order, filename) |
