from typing import Any, List
import re
import os
import ntpath
import torch
import numpy as np
import dataclasses
from . import utils
channelmap = {"Xrotation": "x", "Yrotation": "y", "Zrotation": "z"}
channelmap_inv = {
"x": "Xrotation",
"y": "Yrotation",
"z": "Zrotation",
}
ordermap = {
"x": 0,
"y": 1,
"z": 2,
}
[docs]
@dataclasses.dataclass
class Anim:
"""
Basic animation structure
"""
quats: torch.Tensor
"""Local quaternions"""
pos: torch.Tensor
"""Local positions"""
offsets: torch.Tensor
"""Local joint offsets"""
parents: Any
"""Bone hierarchy"""
bones: List[str]
"""A list of joints"""
frametime: float
"""Frame time slot"""
fps: float = None
"""Frequency"""
def __post_init__(self):
self.fps = 1 / self.frametime
[docs]
def read_bvh(filename: str, start: int = None, end: int = None, order=None) -> Anim:
"""Reading a BVH file and extracts animation information.
Args:
filename (str): File path.
start (int, optional): Strat frame idx. Defaults to None.
end (int, optional): Ending frame idx. Defaults to None.
order (_type_, optional): Order of euler rotations. Defaults to None.
Returns:
Anim: Parsed Animation object
"""
i = 0
active = -1
end_site = False
names = []
orients = np.array([]).reshape((0, 4))
offsets = np.array([]).reshape((0, 3))
parents = np.array([], dtype=int)
# Parse the file, line by line
with open(filename, "r") as f:
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))
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 = 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))
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
rotations = utils.euler_to_quat(np.radians(rotations), order=order)
rotations = utils.remove_quat_discontinuities(rotations)
return Anim(rotations, positions, offsets, parents, names, frametime)
# """
# Reads a BVH file and extracts animation information.
# :param filename: BVh filename
# :param start: start frame
# :param end: end frame
# :param order: order of euler rotations
# :return: A simple Anim object conatining the extracted information.
# """
[docs]
def read_bvh_multi_channel(
filename: str, start: int = None, end: int = None, order: Any = None
) -> Anim:
"""Reading a BVH data that has multi channels.
Args:
filename (str): File path.
start (int, optional): Strating frame index. Defaults to None.
end (int, optional): Ending frame index. Defaults to None.
order (Any, optional): Ordering. Defaults to None.
Returns:
Anim: Animation instance
"""
f = open(filename, "r")
i = 0
active = -1
end_site = False
names = []
orients = np.array([]).reshape((0, 4))
offsets = np.array([]).reshape((0, 3))
parents = np.array([], dtype=int)
channel_start_idx = 0
channel_idx_list = [channel_start_idx]
# Parse the file, line by line
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))
channel_start_idx += channels
channel_idx_list.append(channel_start_idx)
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 = 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))
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
start_idx = 0
for j in range(len(channel_idx_list) - 1):
channel_num = channel_idx_list[j + 1] - channel_idx_list[j]
if channel_num == 3:
rotations[fi, j : j + 1] = data_block[start_idx : start_idx + 3]
start_idx += 3
elif channel_num == 6:
positions[fi, j : j + 1] += data_block[start_idx : start_idx + 3]
rotations[fi, j : j + 1] = data_block[start_idx + 3 : start_idx + 6]
start_idx += 6
else:
raise NotImplementedError
i += 1
f.close()
rotations = utils.euler_to_quat(np.radians(rotations), order=order)
rotations = utils.remove_quat_discontinuities(rotations)
return Anim(rotations, positions, offsets, parents, names, frametime)
[docs]
def get_lafan1_set(bvh_path: str, actors, window=50, offset=20):
"""
Extract the same test set as in the article, given the location of the BVH files.
:param bvh_path: Path to the dataset BVH files
:param list: actor prefixes to use in set
:param window: width of the sliding windows (in timesteps)
:param offset: offset between windows (in timesteps)
:return: tuple:
X: local positions
Q: local quaternions
parents: list of parent indices defining the bone hierarchy
contacts_l: binary tensor of left-foot contacts of shape (Batchsize, Timesteps, 2)
contacts_r: binary tensor of right-foot contacts of shape (Batchsize, Timesteps, 2)
"""
npast = 10
subjects = []
seq_names = []
X = []
Q = []
contacts_l = []
contacts_r = []
# Extract
bvh_files = os.listdir(bvh_path)
for file in bvh_files:
if file.endswith(".bvh"):
seq_name, subject = ntpath.basename(file[:-4]).split("_")
if subject in actors:
print("Processing file {}".format(file))
seq_path = os.path.join(bvh_path, file)
anim = read_bvh(seq_path)
# Sliding windows
i = 0
while i + window < anim.pos.shape[0]:
q, x = utils.quat_fk(
anim.quats[i : i + window], anim.pos[i : i + window], anim.parents
)
# Extract contacts
c_l, c_r = utils.extract_feet_contacts(x, [3, 4], [7, 8], velfactor=0.02)
X.append(anim.pos[i : i + window])
Q.append(anim.quats[i : i + window])
seq_names.append(seq_name)
subjects.append(subjects)
contacts_l.append(c_l)
contacts_r.append(c_r)
i += offset
X = np.asarray(X)
Q = np.asarray(Q)
contacts_l = np.asarray(contacts_l)
contacts_r = np.asarray(contacts_r)
# Sequences around XZ = 0
xzs = np.mean(X[:, :, 0, ::2], axis=1, keepdims=True)
X[:, :, 0, 0] = X[:, :, 0, 0] - xzs[..., 0]
X[:, :, 0, 2] = X[:, :, 0, 2] - xzs[..., 1]
# Unify facing on last seed frame
X, Q = utils.rotate_at_frame(X, Q, anim.parents, n_past=npast)
return X, Q, anim.parents, contacts_l, contacts_r
[docs]
def get_train_stats(bvh_folder, train_set):
"""
Extract the same training set as in the paper in order to compute the normalizing statistics
:return: Tuple of (local position mean vector, local position standard deviation vector, local joint offsets tensor)
"""
print("Building the train set...")
xtrain, qtrain, parents, _, _ = get_lafan1_set(bvh_folder, train_set, window=50, offset=20)
print("Computing stats...\n")
# Joint offsets : are constant, so just take the first frame:
offsets = xtrain[0:1, 0:1, 1:, :] # Shape : (1, 1, J, 3)
# Global representation:
q_glbl, x_glbl = utils.quat_fk(qtrain, xtrain, parents)
# Global positions stats:
x_mean = np.mean(
x_glbl.reshape([x_glbl.shape[0], x_glbl.shape[1], -1]).transpose([0, 2, 1]),
axis=(0, 2),
keepdims=True,
)
x_std = np.std(
x_glbl.reshape([x_glbl.shape[0], x_glbl.shape[1], -1]).transpose([0, 2, 1]),
axis=(0, 2),
keepdims=True,
)
return x_mean, x_std, offsets