Source code for rlightning.humanoid.utils.lafan_vendor.extract

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