""" Copyright (c) Meta Platforms, Inc. and affiliates. All rights reserved. This source code is licensed under the license found in the LICENSE file in the root directory of this source tree. """ import json import numpy as np import re import torch import torch as th import torch.nn as nn import torch.nn.functional as F from typing import Dict, Any from visualize.ca_body.utils.quaternion import Quaternion from pytorch3d.transforms import matrix_to_euler_angles from typing import Optional, Tuple import logging logger = logging.getLogger(__name__) class ParameterTransform(nn.Module): def __init__(self, lbs_cfg_dict: Dict[str, Any]): super().__init__() # self.pose_names = list(lbs_cfg_dict["joint_names"]) self.channel_names = list(lbs_cfg_dict["channel_names"]) transform_offsets = torch.FloatTensor(lbs_cfg_dict["transform_offsets"]) transform = torch.FloatTensor(lbs_cfg_dict["transform"]) self.limits = lbs_cfg_dict["limits"] self.nr_scaling_params = lbs_cfg_dict["nr_scaling_params"] self.nr_position_params = lbs_cfg_dict["nr_position_params"] self.nr_total_params = self.nr_scaling_params + self.nr_position_params self.register_buffer("transform_offsets", transform_offsets) self.register_buffer("transform", transform) def forward(self, pose: th.Tensor) -> th.Tensor: """ :param pose: raw pose inputs, shape (batch_size, len(pose_names)) :return: skeleton parameters, shape (batch_size, len(channel_names)*nr_skeleton_joints) """ return self.transform.mm(pose.t()).t() + self.transform_offsets class LinearBlendSkinning(nn.Module): def __init__( self, model_json: Dict[str, Any], lbs_config_dict: Dict[str, Any], num_max_skin_joints: int =8, scale_path: str =None, ): super().__init__() model = model_json self.param_transform = ParameterTransform(lbs_config_dict) self.joint_names = [] nr_joints = len(model["Skeleton"]["Bones"]) joint_parents = torch.zeros((nr_joints, 1), dtype=torch.int64) joint_rotation = torch.zeros((nr_joints, 4), dtype=torch.float32) joint_offset = torch.zeros((nr_joints, 3), dtype=torch.float32) for idx, bone in enumerate(model["Skeleton"]["Bones"]): self.joint_names.append(bone["Name"]) if bone["Parent"] > nr_joints: joint_parents[idx] = -1 else: joint_parents[idx] = bone["Parent"] joint_rotation[idx, :] = torch.FloatTensor(bone["PreRotation"]) joint_offset[idx, :] = torch.FloatTensor(bone["TranslationOffset"]) skin_model = model["SkinnedModel"] mesh_vertices = torch.FloatTensor(skin_model["RestPositions"]) mesh_normals = torch.FloatTensor(skin_model["RestVertexNormals"]) weights = torch.FloatTensor([e[1] for e in skin_model["SkinningWeights"]]) indices = torch.LongTensor([e[0] for e in skin_model["SkinningWeights"]]) offsets = torch.LongTensor(skin_model["SkinningOffsets"]) nr_vertices = len(offsets) - 1 skin_weights = torch.zeros((nr_vertices, num_max_skin_joints), dtype=torch.float32) skin_indices = torch.zeros((nr_vertices, num_max_skin_joints), dtype=torch.int64) offset_right = offsets[1:] for offset in range(num_max_skin_joints): offset_left = offsets[:-1] + offset skin_weights[offset_left < offset_right, offset] = weights[ offset_left[offset_left < offset_right] ] skin_indices[offset_left < offset_right, offset] = indices[ offset_left[offset_left < offset_right] ] mesh_faces = torch.IntTensor(skin_model["Faces"]["Indices"]).view(-1, 3) mesh_texture_faces = torch.IntTensor(skin_model["Faces"]["TextureIndices"]).view(-1, 3) mesh_texture_coords = torch.FloatTensor(skin_model["TextureCoordinates"]).view(-1, 2) # zero_pose = torch.zeros((1, len(self.param_transform.pose_names)), dtype=torch.float32) zero_pose = torch.zeros((1, self.param_transform.nr_total_params), dtype=torch.float32) bind_state = solve_skeleton_state( self.param_transform(zero_pose), joint_offset, joint_rotation, joint_parents ) # self.register_buffer('mesh_vertices', mesh_vertices) # we want to train on rest pose # self.mesh_vertices = nn.Parameter(mesh_vertices, requires_grad=optimize_mesh) self.register_buffer("mesh_vertices", mesh_vertices) self.register_buffer("joint_parents", joint_parents) self.register_buffer("joint_rotation", joint_rotation) self.register_buffer("joint_offset", joint_offset) self.register_buffer("mesh_normals", mesh_normals) self.register_buffer("mesh_faces", mesh_faces) self.register_buffer("mesh_texture_faces", mesh_texture_faces) self.register_buffer("mesh_texture_coords", mesh_texture_coords) self.register_buffer("skin_weights", skin_weights) self.register_buffer("skin_indices", skin_indices) self.register_buffer("bind_state", bind_state) self.register_buffer("rest_vertices", mesh_vertices) # pre-compute joint weights self.register_buffer("joints_weights", self.compute_joints_weights()) if scale_path is not None: scale = np.loadtxt(scale_path).astype(np.float32)[np.newaxis] scale = scale[:, 0, :] if len(scale.shape) == 3 else scale self.register_buffer("scale", torch.tensor(scale)) @property def num_verts(self): return self.mesh_vertices.size(0) @property def num_joints(self): return self.joint_offset.size(0) @property def num_params(self): return self.skin_weights.shape[-1] def compute_rigid_transforms(self, global_pose: th.Tensor, local_pose: th.Tensor, scale: th.Tensor): """Returns rigid transforms.""" params = torch.cat([global_pose, local_pose, scale], axis=-1) params = self.param_transform(params) return solve_skeleton_state( params, self.joint_offset, self.joint_rotation, self.joint_parents ) def compute_rigid_transforms_matrix(self, global_pose: th.Tensor, local_pose: th.Tensor, scale: th.Tensor): params = torch.cat([global_pose, local_pose, scale], axis=-1) params = self.param_transform(params) states = solve_skeleton_state( params, self.joint_offset, self.joint_rotation, self.joint_parents ) return states_to_matrix(self.bind_state, states) def compute_joints_weights(self, drop_empty=False): """Compute weights per joint given flattened weights-indices.""" idxs_verts = torch.arange(self.num_verts)[:, np.newaxis].expand(-1, self.num_params) weights_joints = torch.zeros( (self.num_joints, self.num_verts), dtype=torch.float32, device=self.skin_weights.device, ) weights_joints[self.skin_indices, idxs_verts] = self.skin_weights if drop_empty: weights_joints = weights_joints[weights_joints.sum(axis=-1).abs() > 0] return weights_joints def compute_root_rigid_transform(self, poses: th.Tensor) -> Tuple[th.Tensor, th.Tensor]: """Get a transform of the root joint.""" scales = torch.zeros( (poses.shape[0], self.nr_total_params - poses.shape[1]), dtype=poses.dtype, device=poses.device, ) params = torch.cat((poses, scales), 1) states = solve_skeleton_state( self.param_transform(params), self.joint_offset, self.joint_rotation, self.joint_parents, ) mat = states_to_matrix(self.bind_state, states) return mat[:, 1, :, 3], mat[:, 1, :, :3] def compute_relative_rigid_transforms(self, global_pose: th.Tensor, local_pose: th.Tensor, scale: th.Tensor): params = torch.cat([global_pose, local_pose, scale], axis=-1) params = self.param_transform(params) batch_size = params.shape[0] joint_offset = self.joint_offset joint_rotation = self.joint_rotation # batch processing for parameters jp = params.view((batch_size, -1, 7)) lt = jp[:, :, 0:3] + joint_offset.unsqueeze(0) lr = Quaternion.batchMul(joint_rotation.unsqueeze(0), Quaternion.batchFromXYZ(jp[:, :, 3:6])) return torch.cat([lt, lr], axis=-1) def skinning(self, bind_state: th.Tensor, vertices: th.Tensor, target_states: th.Tensor): """ Apply skinning to a set of states Args: b/bind_state: 1 x nr_joint x 8 bind state v/vertices: 1 x nr_vertices x 3 vertices t/target_states: batch_size x nr_joint x 8 current states Returns: batch_size x nr_vertices x 3 skinned vertices """ assert target_states.size()[1:] == bind_state.size()[1:] mat = states_to_matrix(bind_state, target_states) # apply skinning to vertices vs = torch.matmul( mat[:, self.skin_indices], torch.cat((vertices, torch.ones_like(vertices[:, :, 0]).unsqueeze(2)), dim=2) .unsqueeze(2) .unsqueeze(4), ) ws = self.skin_weights.unsqueeze(2).unsqueeze(3) res = (vs * ws).sum(dim=2).squeeze(3) return res def unpose(self, poses: th.Tensor, scales: th.Tensor, verts: th.Tensor): """ :param poses: 100 (tx ty tz rx ry rz) params in blueman :param scales: 29 (s) params in blueman :return: """ # check shape of poses and scales params = torch.cat((poses, scales), 1) states = solve_skeleton_state( self.param_transform(params), self.joint_offset, self.joint_rotation, self.joint_parents, ) return self.unskinning(self.bind_state, states, verts) def unskinning(self, bind_state: th.Tensor, target_states: th.Tensor, verts: th.Tensor): """Apply skinning to a set of states Args: bind_state: [B, NJ, 8] - bind state target_states: [B, NJ, 8] - current states vertices: [B, V, 3] - vertices Returns: batch_size x nr_vertices x 3 skinned vertices """ assert target_states.size()[1:] == bind_state.size()[1:] mat = states_to_matrix(bind_state, target_states) ws = self.skin_weights[None, :, :, None, None] sum_mat = (mat[:, self.skin_indices] * ws).sum(dim=2) sum_mat4x4 = torch.cat((sum_mat, torch.zeros_like(sum_mat[:, :, :1, :])), dim=2) sum_mat4x4[:, :, 3, 3] = 1.0 verts_4d = torch.cat((verts, torch.ones_like(verts[:, :, :1])), dim=2).unsqueeze(3) resmesh = [] for i in range(sum_mat.shape[0]): newmat = sum_mat4x4[i, :, :, :].contiguous() invnewmat = newmat.inverse() tmpvets = invnewmat.matmul(verts_4d[i]) resmesh.append(tmpvets.unsqueeze(0)) resmesh = torch.cat(resmesh) return resmesh.squeeze(3)[..., :3].contiguous() def forward(self, poses: th.Tensor, scales: th.Tensor, verts_unposed: Optional[th.Tensor] = None) -> th.Tensor: """ Args: poses: [B, NP] - pose parametersa scales: [B, NS] - additional scaling params verts_unposed: [B, N, 3] - unposed vertices Returns: [B, N, 3] - posed vertices """ params = torch.cat((poses, scales), 1) params_transformed = self.param_transform(params) states = solve_skeleton_state( params_transformed, self.joint_offset, self.joint_rotation, self.joint_parents, ) if verts_unposed is None: mesh = self.skinning(self.bind_state, self.mesh_vertices.unsqueeze(0), states) else: mesh = self.skinning(self.bind_state, verts_unposed, states) return mesh def solve_skeleton_state(param: th.Tensor, joint_offset: th.Tensor, joint_rotation: th.Tensor, joint_parents: th.Tensor): """ :param param: batch_size x (7*nr_skeleton_joints) ParamTransform Outputs. :return: batch_size x nr_skeleton_joints x 8 Skeleton States 8 stands form 3 translation + 4 rotation (quat) + 1 scale """ batch_size = param.shape[0] # batch processing for parameters jp = param.view((batch_size, -1, 7)) lt = jp[:, :, 0:3] + joint_offset.unsqueeze(0) lr = Quaternion.batchMul(joint_rotation.unsqueeze(0), Quaternion.batchFromXYZ(jp[:, :, 3:6])) ls = torch.pow( torch.tensor([2.0], dtype=torch.float32, device=param.device), jp[:, :, 6].unsqueeze(2), ) state = [] for index, parent in enumerate(joint_parents): if int(parent) != -1: gr = Quaternion.batchMul(state[parent][:, :, 3:7], lr[:, index, :].unsqueeze(1)) gt = ( Quaternion.batchRot( state[parent][:, :, 3:7], lt[:, index, :].unsqueeze(1) * state[parent][:, :, 7].unsqueeze(2), ) + state[parent][:, :, 0:3] ) gs = state[parent][:, :, 7].unsqueeze(2) * ls[:, index, :].unsqueeze(1) state.append(torch.cat((gt, gr, gs), dim=2)) else: state.append( torch.cat((lt[:, index, :], lr[:, index, :], ls[:, index, :]), dim=1).view( (batch_size, 1, 8) ) ) return torch.cat(state, dim=1) def states_to_matrix(bind_state: th.Tensor, target_states: th.Tensor, return_transform: bool=False): # multiply bind inverse with states br = Quaternion.batchInvert(bind_state[:, :, 3:7]) bs = bind_state[:, :, 7].unsqueeze(2).reciprocal() bt = Quaternion.batchRot(br, -bind_state[:, :, 0:3]) * bs # applying rotation tr = Quaternion.batchMul(target_states[:, :, 3:7], br) # applying scaling ts = target_states[:, :, 7].unsqueeze(2) * bs # applying transformation tt = ( Quaternion.batchRot(target_states[:, :, 3:7], bt * target_states[:, :, 7].unsqueeze(2)) + target_states[:, :, 0:3] ) # convert to matrices twx = 2.0 * tr[:, :, 0] * tr[:, :, 3] twy = 2.0 * tr[:, :, 1] * tr[:, :, 3] twz = 2.0 * tr[:, :, 2] * tr[:, :, 3] txx = 2.0 * tr[:, :, 0] * tr[:, :, 0] txy = 2.0 * tr[:, :, 1] * tr[:, :, 0] txz = 2.0 * tr[:, :, 2] * tr[:, :, 0] tyy = 2.0 * tr[:, :, 1] * tr[:, :, 1] tyz = 2.0 * tr[:, :, 2] * tr[:, :, 1] tzz = 2.0 * tr[:, :, 2] * tr[:, :, 2] mat = torch.stack( ( torch.stack((1.0 - (tyy + tzz), txy + twz, txz - twy), dim=2) * ts, torch.stack((txy - twz, 1.0 - (txx + tzz), tyz + twx), dim=2) * ts, torch.stack((txz + twy, tyz - twx, 1.0 - (txx + tyy)), dim=2) * ts, tt, ), dim=3, ) if return_transform: return mat, (tr, tt, ts) return mat def get_influence_map( transform_raw: th.Tensor, pose_length=None, num_params_per_joint=7, eps=1.0e-6 ): num_joints = transform_raw.shape[0] // num_params_per_joint num_params = transform_raw.shape[-1] if pose_length is None: pose_length = num_params assert pose_length <= num_params transform_raw = transform_raw.reshape((num_joints, num_params_per_joint, num_params)) return [ torch.where(torch.abs(transform_raw[i, :, :pose_length]) > eps)[1].tolist() for i in range(num_joints) ] def compute_weights_joints_slow(lbs_weights, lbs_indices, num_joints): num_verts = lbs_weights.shape[0] weights_joints = torch.zeros((num_joints, num_verts), dtype=torch.float32) for i in range(num_verts): idx = lbs_indices[i, :] weights_joints[idx, i] = lbs_weights[i, :] return weights_joints def load_momentum_cfg(model, lbs_config_txt_fh, nr_scaling_params=None): def find(l, x): try: return l.index(x) except ValueError: return None """Load a parameter configuration file""" channelNames = ["tx", "ty", "tz", "rx", "ry", "rz", "sc"] paramNames = [] joint_names = [] for idx, bone in enumerate(model["Skeleton"]["Bones"]): joint_names.append(bone["Name"]) def findJointIndex(x): return find(joint_names, x) def findParameterIndex(x): return find(paramNames, x) limits = [] # create empty result transform_triplets = [] lines = lbs_config_txt_fh.readlines() # read until end for line in lines: # strip comments line = line[: line.find("#")] if line.find("limit") != -1: r = re.search("limit ([\\w.]+) (\\w+) (.*)", line) if r is None: continue if len(r.groups()) != 3: logger.info("Failed to parse limit configuration line :\n " + line) continue # find parameter and/or joint index fullname = r.groups()[0] type = r.groups()[1] remaining = r.groups()[2] parameterIndex = findParameterIndex(fullname) jointName = fullname.split(".") jointIndex = findJointIndex(jointName[0]) channelIndex = -1 if jointIndex is not None and len(jointName) == 2: # find matching channel name channelIndex = channelNames.index(jointName[1]) if channelIndex is None: logger.info( "Unknown joint channel name " + jointName[1] + " in parameter configuration line :\n " + line ) continue # only parse passive limits for now if type == "minmax_passive" or type == "minmax": # match [ , ] rp = re.search( "\\[\\s*([-+]?[0-9]*\\.?[0-9]+)\\s*,\\s*([-+]?[0-9]*\\.?[0-9]+)\\s*\\](\\s*[-+]?[0-9]*\\.?[0-9]+)?", remaining, ) if len(rp.groups()) != 3: logger.info(f"Failed to parse passive limit configuration line :\n {line}") continue minVal = float(rp.groups()[0]) maxVal = float(rp.groups()[1]) weightVal = 1.0 if len(rp.groups()) == 3 and not rp.groups()[2] is None: weightVal = float(rp.groups()[2]) # result.limits.append([jointIndex * 7 + channelIndex, minVal, maxVal]) if channelIndex >= 0: valueIndex = jointIndex * 7 + channelIndex limit = { "type": "LimitMinMaxJointValue", "str": fullname, "valueIndex": valueIndex, "limits": [minVal, maxVal], "weight": weightVal, } limits.append(limit) else: if parameterIndex is None: logger.info(f"Unknown parameterIndex : {fullname}\n {line} {paramNames} ") continue limit = { "type": "LimitMinMaxParameter", "str": fullname, "parameterIndex": parameterIndex, "limits": [minVal, maxVal], "weight": weightVal, } limits.append(limit) # continue the remaining file continue # check for parameterset definitions and ignore if line.find("parameterset") != -1: continue # use regex to parse definition r = re.search("(\w+).(\w+)\s*=\s*(.*)", line) if r is None: continue if len(r.groups()) != 3: logger.info("Failed to parse parameter configuration line :\n " + line) continue # find joint name and parameter jointIndex = findJointIndex(r.groups()[0]) if jointIndex is None: logger.info( "Unknown joint name " + r.groups()[0] + " in parameter configuration line :\n " + line ) continue # find matching channel name channelIndex = channelNames.index(r.groups()[1]) if channelIndex is None: logger.info( "Unknown joint channel name " + r.groups()[1] + " in parameter configuration line :\n " + line ) continue valueIndex = jointIndex * 7 + channelIndex # parse parameters parameterList = r.groups()[2].split("+") for parameterPair in parameterList: parameterPair = parameterPair.strip() r = re.search("\s*([+-]?[0-9]*\.?[0-9]*)\s\*\s(\w+)\s*", parameterPair) if r is None or len(r.groups()) != 2: logger.info( "Malformed parameter description " + parameterPair + " in parameter configuration line :\n " + line ) continue val = float(r.groups()[0]) parameter = r.groups()[1] # check if parameter exists parameterIndex = findParameterIndex(parameter) if parameterIndex is None: # no, create new parameter entry parameterIndex = len(paramNames) paramNames.append(parameter) transform_triplets.append((valueIndex, parameterIndex, val)) # set (dense) parameter_transformation matrix transform = np.zeros((len(channelNames) * len(joint_names), len(paramNames)), dtype=np.float32) for i, j, v in transform_triplets: transform[i, j] = v outputs = { "model_param_names": paramNames, "joint_names": joint_names, "channel_names": channelNames, "limits": limits, "transform": transform, "transform_offsets": np.zeros((1, len(channelNames) * len(joint_names)), dtype=np.float32), } # set number of scales automatically if nr_scaling_params is None: outputs.update(nr_scaling_params=len([s for s in paramNames if s.startswith("scale")])) outputs.update(nr_position_params=len(paramNames) - outputs["nr_scaling_params"]) return outputs def compute_normalized_pose_quat(lbs, local_pose, scale): """Computes a normalized representation of the pose in quaternion space. This is a delta between the per-joint local transformation and the bind state. Returns: [B, NJ, 4] - normalized rotations """ B = local_pose.shape[0] global_pose_zero = th.zeros((B, 6), dtype=th.float32, device=local_pose.device) params = lbs.param_transform(th.cat([global_pose_zero, local_pose, scale], axis=-1)) params = params.reshape(B, -1, 7) # applying rotation # TODO: what is this? rot_quat = Quaternion.batchMul(lbs.joint_rotation[np.newaxis], Quaternion.batchFromXYZ(params[:, :, 3:6])) # removing the bind state bind_rot_quat = Quaternion.batchInvert(lbs.bind_state[:, :, 3:7]) return Quaternion.batchMul(rot_quat, bind_rot_quat) def compute_root_transform_cuda(lbs_fn, poses, verts=None): # NOTE: verts is not really necessary, # NOTE: should be used in conjuncation with LBSCuda B = poses.shape[0] # NOTE: scales are zero (!) _, _, _, state_t, state_r, state_s = lbs_fn(poses, vertices=verts) bind_r = lbs_fn.joint_state_r_zero[np.newaxis, 1].expand(B, -1, -1) bind_t = lbs_fn.joint_state_t_zero[np.newaxis, 1].expand(B, -1) R_root = th.matmul(state_r[:, 1], bind_r) t_root = ( th.matmul(state_r[:, 1], (bind_t * state_s[:, 1])[..., np.newaxis])[..., 0] + state_t[:, 1] ) return R_root, t_root # def compute_joints_weights(lbs_fn: LinearBlendSkinningCuda, drop_empty: bool = False) -> th.Tensor: # device = lbs_fn.skin_indices.device # idxs_verts = th.arange(lbs_fn.nr_vertices)[:, np.newaxis].to(device) # weights_joints = th.zeros( # (lbs_fn.nr_joints, lbs_fn.nr_vertices), # dtype=th.float32, # device=lbs_fn.skin_indices.device, # ) # weights_joints[lbs_fn.skin_indices, idxs_verts] = lbs_fn.skin_weights # if drop_empty: # weights_joints = weights_joints[weights_joints.sum(axis=-1).abs() > 0] # return weights_joints # def compute_pose_regions(lbs_fn: LinearBlendSkinningCuda) -> np.ndarray: # """Computes pose regions given a linear blend skinning function. # Returns: # np.ndarray of boolean masks of shape [nr_params, n_rvertices] # """ # weights = compute_joints_weights(lbs_fn).cpu().numpy() # n_pos = lbs_fn.nr_position_params # param_masks = np.zeros((n_pos, weights.shape[-1])) # children = {j: [] for j in range(lbs_fn.nr_joints)} # parents = {j: None for j in range(lbs_fn.nr_joints)} # prec = {j: [] for j in range(lbs_fn.nr_joints)} # for j in range(lbs_fn.nr_joints): # parent_index = int(lbs_fn.joint_parents[j]) # if parent_index == -1: # continue # children[parent_index].append(j) # parents[j] = parent_index # prec[j] = [parent_index, int(lbs_fn.joint_parents[parent_index])] # # get parameters for each joint # # j_to_p = get_influence_map(lbs_fn.param_transform.transform, n_pos) # j_to_p = get_influence_map(lbs_fn.param_transform, n_pos) # # get all the joints # p_to_j = [[] for i in range(n_pos)] # for j, pidx in enumerate(j_to_p): # for p in pidx: # if j not in p_to_j[p]: # p_to_j[p].append(j) # for p, jidx in enumerate(p_to_j): # param_masks[p] = weights[jidx].sum(axis=0) # if not np.any(param_masks[p]): # assert len(jidx) == 1 # jidx_c = children[jidx[0]][:] # for jc in jidx_c[:]: # jidx_c += children[jc] # param_masks[p] = weights[jidx_c].sum(axis=0) # return param_masks > 0.0 def compute_pose_regions_legacy(lbs_fn) -> np.ndarray: """Computes pose regions given a linear blend skinning function.""" weights = lbs_fn.joints_weights.cpu().numpy() n_pos = lbs_fn.param_transform.nr_position_params param_masks = np.zeros((n_pos, lbs_fn.joints_weights.shape[-1])) children = {j: [] for j in range(lbs_fn.num_joints)} parents = {j: None for j in range(lbs_fn.num_joints)} prec = {j: [] for j in range(lbs_fn.num_joints)} for j in range(lbs_fn.num_joints): parent_index = int(lbs_fn.joint_parents[j, 0]) if parent_index == -1: continue children[parent_index].append(j) parents[j] = parent_index prec[j] = [parent_index, int(lbs_fn.joint_parents[parent_index, 0])] # get parameters for each joint j_to_p = get_influence_map(lbs_fn.param_transform.transform, n_pos) # get all the joints p_to_j = [[] for i in range(n_pos)] for j, pidx in enumerate(j_to_p): for p in pidx: if j not in p_to_j[p]: p_to_j[p].append(j) for p, jidx in enumerate(p_to_j): param_masks[p] = weights[jidx].sum(axis=0) if not np.any(param_masks[p]): assert len(jidx) == 1 jidx_c = children[jidx[0]][:] for jc in jidx_c[:]: jidx_c += children[jc] param_masks[p] = weights[jidx_c].sum(axis=0) return param_masks > 0.0 def compute_pose_mask_uv(lbs_fn, geo_fn, uv_size, ksize=25): device = geo_fn.index_image.device pose_regions = compute_pose_regions(lbs_fn) pose_regions = ( th.as_tensor(pose_regions[6:], dtype=th.float32).permute(1, 0)[np.newaxis].to(device) ) pose_regions_uv = geo_fn.to_uv(pose_regions) pose_regions_uv = F.max_pool2d(pose_regions_uv, ksize, 1, padding=ksize // 2) pose_cond_mask = (F.interpolate(pose_regions_uv, size=(uv_size, uv_size)) > 0.1).to(th.int32) return pose_cond_mask def parent_chain(joint_parents, idx, depth): if depth == 0 or idx == 0: return [] parent_idx = int(joint_parents[idx]) return [parent_idx] + parent_chain(joint_parents, parent_idx, depth - 1) def joint_connectivity(nr_joints, joint_parents, chain_depth=2, pad_ancestors=False): children = {j: [] for j in range(nr_joints)} parents = {j: None for j in range(nr_joints)} ancestors = {j: [] for j in range(nr_joints)} for j in range(nr_joints): parent_index = int(joint_parents[j]) ancestors[j] = parent_chain(joint_parents, j, depth=chain_depth) if pad_ancestors: # adding itself ancestors[j] += [j] * (chain_depth - len(ancestors[j])) if parent_index == -1: continue children[parent_index].append(j) parents[j] = parent_index return { 'children': children, 'parents': parents, 'ancestors': ancestors, } # TODO: merge this with LinearBlendSkinning? class LBSModule(nn.Module): def __init__( self, lbs_model_json, lbs_config_dict, lbs_template_verts, lbs_scale, global_scaling ): super().__init__() self.lbs_fn = LinearBlendSkinning(lbs_model_json, lbs_config_dict) self.register_buffer("lbs_scale", th.as_tensor(lbs_scale, dtype=th.float32)) self.register_buffer( "lbs_template_verts", th.as_tensor(lbs_template_verts, dtype=th.float32) ) self.register_buffer("global_scaling", th.as_tensor(global_scaling)) def pose(self, verts_unposed, motion, template: Optional[th.Tensor] = None): scale = self.lbs_scale.expand(motion.shape[0], -1) if template is None: template = self.lbs_template_verts return self.lbs_fn(motion, scale, verts_unposed + template) * self.global_scaling def unpose(self, verts, motion): B = motion.shape[0] scale = self.lbs_scale.expand(B, -1) return ( self.lbs_fn.unpose(motion, scale, verts / self.global_scaling) - self.lbs_template_verts ) def template_pose(self, motion): B = motion.shape[0] scale = self.lbs_scale.expand(B, -1) verts = self.lbs_template_verts[np.newaxis].expand(B, -1, -1) return self.lbs_fn(motion, scale, verts) * self.global_scaling[np.newaxis]