# Copyright (c) Meta Platforms, Inc. and affiliates. # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. import logging from animated_drawings.model.bvh import BVH import numpy as np import numpy.typing as npt import math from animated_drawings.model.joint import Joint from sklearn.decomposition import PCA from typing import Tuple, List, Dict from animated_drawings.model.vectors import Vectors from animated_drawings.model.quaternions import Quaternions from animated_drawings.config import MotionConfig, RetargetConfig x_axis = np.array([1.0, 0.0, 0.0], dtype=np.float32) z_axis = np.array([0.0, 0.0, 1.0], dtype=np.float32) class Retargeter(): """ Retargeter class takes in a motion_cfg file and retarget_cfg file. Using the specifications listed within retarget_cfg, it converts the motion specified in motion_cfg into a formal that can be applied to an animated drawing. It is responsible for project 3D joint locations onto 2D planes, determining resulting bone orientations, joint 'depths', and root offsets for each frame. """ def __init__(self, motion_cfg: MotionConfig, retarget_cfg: RetargetConfig) -> None: # instantiate the bvh try: self.bvh = BVH.from_file(str(motion_cfg.bvh_p), motion_cfg.start_frame_idx, motion_cfg.end_frame_idx) except Exception as e: msg = f'Error loading BVH: {e}' logging.critical(msg) assert False, msg # get and cache bvh joint names for later self.bvh_joint_names = self.bvh.get_joint_names() # bvh joints defining a set of vectors that skeleton's fwd is perpendicular to self.forward_perp_vector_joint_names: List[Tuple[str, str]] = motion_cfg.forward_perp_joint_vectors # override the frame_time, if one was specified within motion_cfg if motion_cfg.frame_time: self.bvh.frame_time = motion_cfg.frame_time # rotate BVH skeleton so up is +Y if motion_cfg.up == '+y': pass # no rotation needed elif motion_cfg.up == '+z': self.bvh.set_rotation(Quaternions.from_euler_angles('yx', np.array([-90.0, -90.0]))) else: msg = f'up value not implemented: {motion_cfg.up}' logging.critical(msg) assert False, msg # rotate BVH skeleton so forward is +X skeleton_fwd: Vectors = self.bvh.get_skeleton_fwd(self.forward_perp_vector_joint_names) q: Quaternions = Quaternions.rotate_between_vectors(skeleton_fwd, Vectors([1.0, 0.0, 0.0])) self.bvh.rotation_offset(q) # scale BVH self.bvh.set_scale(motion_cfg.scale) # position above origin self.bvh.offset(-self.bvh.root_joint.get_world_position()) # adjust bvh skeleton y pos by getting groundplane joint... try: groundplane_joint = self.bvh.root_joint.get_transform_by_name(motion_cfg.groundplane_joint) assert isinstance(groundplane_joint, Joint), f'could not find joint by name: {motion_cfg.groundplane_joint}' except Exception as e: msg = f'Error getting groundplane joint: {e}' logging.warning(msg) assert False # ... and moving the bvh so it is on the y=0 plane bvh_groundplane_y = groundplane_joint.get_world_position()[1] self.bvh.offset(np.array([0, -bvh_groundplane_y, 0])) self.joint_positions: npt.NDArray[np.float32] self.fwd_vectors: npt.NDArray[np.float32] self.bvh_root_positions: npt.NDArray[np.float32] self._compute_normalized_joint_positions_and_fwd_vectors() # cache the starting worldspace location of character's root joint self.character_start_loc: npt.NDArray[np.float32] = np.array(retarget_cfg.char_start_loc, dtype=np.float32) # holds world coordinates of character root joint after retargeting self.char_root_positions: npt.NDArray[np.float32] # get & save projection planes self.joint_group_name_to_projection_plane: Dict[ str, npt.NDArray[np.float32]] = {} self.joint_to_projection_plane: Dict[ str, npt.NDArray[np.float32]] = {} for joint_projection_group in retarget_cfg.bvh_projection_bodypart_groups: group_name = joint_projection_group['name'] joint_names = joint_projection_group['bvh_joint_names'] projection_method = joint_projection_group['method'] projection_plane = self._determine_projection_plane_normal(group_name, joint_names, projection_method) self.joint_group_name_to_projection_plane[joint_projection_group['name']] = projection_plane for joint_name in joint_projection_group['bvh_joint_names']: self.joint_to_projection_plane[joint_name] = projection_plane # map character joint names to its orientations self.char_joint_to_orientation: Dict[str, npt.NDArray[np.float32]] = {} # map bvh joint names to its distance to project plane (useful for rendering order) self.bvh_joint_to_projection_depth: Dict[str, npt.NDArray[np.float32]] = self._compute_depths() def _compute_normalized_joint_positions_and_fwd_vectors(self) -> None: """ Called during initialization. Computes fwd vector for bvh skeleton at each frame. Extracts all bvh skeleton joint locations for all frames. Repositions them so root is above the origin. Rotates them so skeleton faces along the +X axis. """ # get joint positions and forward vectors self.joint_positions = np.empty([self.bvh.frame_max_num, 3 * self.bvh.joint_num], dtype=np.float32) self.fwd_vectors = np.empty([self.bvh.frame_max_num, 3], dtype=np.float32) for frame_idx in range(self.bvh.frame_max_num): self.bvh.apply_frame(frame_idx) self.joint_positions[frame_idx] = self.bvh.root_joint.get_chain_worldspace_positions() self.fwd_vectors[frame_idx] = self.bvh.get_skeleton_fwd(self.forward_perp_vector_joint_names).vs[0] # reposition over origin self.bvh_root_positions = self.joint_positions[:, :3] self.joint_positions = self.joint_positions - np.tile(self.bvh_root_positions, [1, len(self.bvh_joint_names)]) # compute angle between skeleton's forward vector and x axis v1 = np.tile(np.array([1.0, 0.0], dtype=np.float32), reps=(self.joint_positions.shape[0], 1)) v2 = self.fwd_vectors dot: npt.NDArray[np.float32] = v1[:, 0]*v2[:, 0] + v1[:, 1]*v2[:, 2] det: npt.NDArray[np.float32] = v1[:, 0]*v2[:, 2] - v2[:, 0]*v1[:, 1] angle: npt.NDArray[np.float32] = np.arctan2(det, dot).astype(np.float32) angle %= 2*np.pi angle = np.where(angle < 0.0, angle + 2*np.pi, angle) # rotate the skeleton's joint so it faces +X axis for idx in range(self.joint_positions.shape[0]): rot_mat = np.identity(3).astype(np.float32) rot_mat[0, 0] = math.cos(angle[idx]) rot_mat[0, 2] = math.sin(angle[idx]) rot_mat[2, 0] = -math.sin(angle[idx]) rot_mat[2, 2] = math.cos(angle[idx]) rotated_joints: npt.NDArray[np.float32] = rot_mat @ self.joint_positions[idx].reshape([-1, 3]).T self.joint_positions[idx] = rotated_joints.T.reshape(self.joint_positions[idx].shape) def _determine_projection_plane_normal(self, group_name: str, joint_names: List[str], projection_method: str) -> npt.NDArray[np.float32]: """ Given a joint_projection_group dictionary object, computes the projection plane normal used for the group. Called during initialization. """ if projection_method == 'frontal': logging.info(f'{group_name} projection_method is {projection_method}. Using {x_axis}') return x_axis elif projection_method == 'sagittal': logging.info(f'{group_name} projection_method is {projection_method}. Using {z_axis}') return z_axis elif projection_method == 'pca': logging.info(f'{group_name} projection_method is {projection_method}. Running PCA on {joint_names}') pass # pca code is below else: msg = f'bad project method for {group_name}: {projection_method}' logging.critical(msg) assert False, msg # get the xyz locations of joints within this joint_projection_group joints_idxs = [self.bvh_joint_names.index(joint_name) for joint_name in joint_names] joints_mask = np.full(self.joint_positions.shape[1], False, dtype=np.bool8) for idx in joints_idxs: joints_mask[3*idx:3*(idx+1)] = True joints_points = self.joint_positions[:, joints_mask] joints_points = joints_points.reshape([-1, 3]) # do PCA and get 3rd component pca = PCA() pca.fit(joints_points) pc3: npt.NDArray[np.float32] = pca.components_[2] # pyright: ignore[reportGeneralTypeIssues] # see if it is closer to the x axis or z axis x_cos_sim: float = np.dot(x_axis, pc3) / (np.linalg.norm(x_axis) * np.linalg.norm(pc3)) z_cos_sim: float = np.dot(z_axis, pc3) / (np.linalg.norm(z_axis) * np.linalg.norm(pc3)) # return close of the two if abs(x_cos_sim) > abs(z_cos_sim): logging.info(f'PCA complete. {group_name} using {x_axis}') return x_axis else: logging.info(f'PCA complete. {group_name} using {z_axis}') return z_axis def _compute_depths(self) -> Dict[str, npt.NDArray[np.float32]]: """ For each BVH joint within bvh_projection_mapping_groups, compute distance to projection plane. This distance used if the joint is a char_body_segmentation_groups depth_driver. """ bvh_joint_to_projection_depth: Dict[str, npt.NDArray[np.float32]] = {} for joint_name in self.bvh_joint_names: joint_idx = self.bvh_joint_names.index(joint_name) joint_xyz = self.joint_positions[:, 3*joint_idx:3*(joint_idx+1)] try: projection_plane_normal = self.joint_to_projection_plane[joint_name] except Exception: msg = f' error finding projection plane for joint_name: {joint_name}' logging.info(msg) continue # project bone onto 2D plane if np.array_equal(projection_plane_normal, x_axis): joint_depths = joint_xyz[:, 0] elif np.array_equal(projection_plane_normal, z_axis): joint_depths = joint_xyz[:, 2] else: msg = 'error projection_plane_normal' logging.critical(msg) assert False, msg bvh_joint_to_projection_depth[joint_name] = joint_depths return bvh_joint_to_projection_depth def scale_root_positions_for_character(self, char_to_bvh_scale: float, projection_bodypart_group_for_offset: str) -> None: """ Uses projection plane of projection_bodypart_group_for_offset to determine bvh skeleton's projected root offset. Scales that offset to account for differences in lengths of character and bvh skeleton limbs. """ try: projection_plane = self.joint_group_name_to_projection_plane[projection_bodypart_group_for_offset] except Exception as e: msg = f'Error getting projection plane: {str(e)}' logging.critical(msg) assert False, msg self.char_root_positions = np.empty([self.bvh_root_positions.shape[0], 2], dtype=np.float32) self.char_root_positions[0] = [0, 0] for idx in range(1, self.bvh_root_positions.shape[0]): if np.array_equal(projection_plane, np.array([0.0, 0.0, 1.0])): # if sagittal projection v1 = self.fwd_vectors[idx] # we're interested in forward motion else: # if frontal projection v1 = self.fwd_vectors[idx][::-1]*np.array([-1, 1, -1]) # we're interested in lateral motion delta = self.bvh_root_positions[idx] - self.bvh_root_positions[idx-1] # scale root delta for both x and y offsets. Project onto v1 for x offset self.char_root_positions[idx, 0] = self.char_root_positions[idx-1, 0] + char_to_bvh_scale * np.dot(v1, delta) # x self.char_root_positions[idx, 1] = self.char_root_positions[idx-1, 1] + char_to_bvh_scale * delta[1] # y def compute_orientations(self, bvh_prox_joint_name: str, bvh_dist_joint_name: str, char_joint_name: str) -> None: """ Calculates the orientation (degrees CCW of +Y axis) of the vector from bvh_prox_joint->bvh_dist_joint using the projection plane of bvh_dist_joint. Results are saved into a dictionary using char_joint_name as the key. """ # get distal end joint dist_joint = self.bvh.root_joint.get_transform_by_name(bvh_dist_joint_name) if dist_joint is None or not isinstance(dist_joint, Joint) or dist_joint.name is None: msg = 'error finding joint {bvh_dist_joint_name}' logging.critical(msg) assert False, msg # get prox joint prox_joint = self.bvh.root_joint.get_transform_by_name(bvh_prox_joint_name) if prox_joint is None or not isinstance(prox_joint, Joint) or prox_joint.name is None: msg = 'joint {bvh_prox_joint_name} has no parent joint, therefore no bone orientation. Returning zero' logging.info(msg) self.char_joint_to_orientation[char_joint_name] = np.zeros(self.joint_positions.shape[0], dtype=np.float32) return # get joint xyz locations dist_joint_idx = self.bvh_joint_names.index(dist_joint.name) dist_joint_xyz = self.joint_positions[:, 3*dist_joint_idx:3*(dist_joint_idx+1)] prox_joint_idx = self.bvh_joint_names.index(prox_joint.name) prox_joint_xyz = self.joint_positions[:, 3*prox_joint_idx:3*(prox_joint_idx+1)] # compute the bone vector bone_vector = dist_joint_xyz - prox_joint_xyz # type: ignore # get distal joint's projection plane try: projection_plane_normal = self.joint_to_projection_plane[bvh_dist_joint_name] except Exception: msg = f' error finding projection plane for bvh_end_joint_name: {bvh_dist_joint_name}' logging.critical(msg) assert False, msg # project bone onto 2D plane if np.array_equal(projection_plane_normal, x_axis): projected_bone_xy = np.stack((-bone_vector[:, 2], bone_vector[:, 1]), axis=1) elif np.array_equal(projection_plane_normal, z_axis): projected_bone_xy = np.stack((bone_vector[:, 0], bone_vector[:, 1]), axis=1) else: msg = 'error projection_plane_normal' logging.critical(msg) assert False, msg # get angles between y axis and bone projected_bone_xy /= np.expand_dims(np.linalg.norm(projected_bone_xy, axis=1), axis=-1) # norm vector y_axis = np.tile(np.array([0.0, 1.0]), reps=(projected_bone_xy.shape[0], 1)) at1 = np.arctan2(projected_bone_xy[:, 1], projected_bone_xy[:, 0], dtype=np.float32) at2 = np.arctan2(y_axis[:, 1], y_axis[:, 0], dtype=np.float32) theta: npt.NDArray[np.float32] = at1 - at2 # type: ignore theta = np.degrees(theta) % 360.0 theta = np.where(theta < 0.0, theta + 360, theta) # save it self.char_joint_to_orientation[char_joint_name] = np.array(theta) def get_retargeted_frame_data(self, time: float) -> Tuple[Dict[str, float], Dict[str, float], npt.NDArray[np.float32]]: """ Input: time, in seconds, used to select the correct BVH frame. Calculate the proper frame and, for it, returns: - orientations, dictionary mapping from character joint names to world orientations (degrees CCW from +Y axis) - joint_depths, dictionary mapping from BVH skeleton's joint names to distance from joint to projection plane - root_positions, the position of the character's root at this frame. """ frame_idx = int(round(time / self.bvh.frame_time, 0)) if frame_idx < 0: logging.info(f'invalid frame_idx ({frame_idx}), replacing with 0') frame_idx = 0 if self.bvh.frame_max_num <= frame_idx: logging.info(f'invalid frame_idx ({frame_idx}), replacing with last frame {self.bvh.frame_max_num-1}') frame_idx = self.bvh.frame_max_num-1 orientations = {key: val[frame_idx] for (key, val) in self.char_joint_to_orientation.items()} joint_depths = {key: val[frame_idx] for (key, val) in self.bvh_joint_to_projection_depth.items()} root_position = np.array([self.char_root_positions[frame_idx, 0], self.char_root_positions[frame_idx, 1], 0.0], dtype=np.float32) root_position += self.character_start_loc # offset by character's starting location return orientations, joint_depths, root_position