def predict_skeleton(input_data, vox, root_pred_net, bone_pred_net, mesh_filename): """ Predict skeleton structure based on joints :param input_data: wrapped data :param vox: voxelized mesh :param root_pred_net: network to predict root :param bone_pred_net: network to predict pairwise connectivity cost :param mesh_filename: meshfilename for debugging :return: predicted skeleton structure """ root_id = getInitId(input_data, root_pred_net) pred_joints = input_data.y[:input_data.num_joint[0]].data.cpu().numpy() with torch.no_grad(): connect_prob, _ = bone_pred_net(input_data) connect_prob = torch.sigmoid(connect_prob) pair_idx = input_data.pairs.long().data.cpu().numpy() prob_matrix = np.zeros((data.num_joint[0], data.num_joint[0])) prob_matrix[pair_idx[:, 0], pair_idx[:, 1]] = connect_prob.data.cpu().numpy().squeeze() prob_matrix = prob_matrix + prob_matrix.transpose() cost_matrix = -np.log(prob_matrix + 1e-10) cost_matrix = increase_cost_for_outside_bone(cost_matrix, pred_joints, vox) pred_skel = Info() parent, key = primMST_symmetry(cost_matrix, root_id, pred_joints) for i in range(len(parent)): if parent[i] == -1: pred_skel.root = TreeNode('root', tuple(pred_joints[i])) break loadSkel_recur(pred_skel.root, i, None, pred_joints, parent) pred_skel.joint_pos = pred_skel.get_joint_dict() #show_mesh_vox(mesh_filename, vox, pred_skel.root) img = show_obj_skel(mesh_filename, pred_skel.root) return pred_skel
def tranfer_to_ori_mesh(filename_ori, filename_remesh, pred_rig): """ convert the predicted rig of remeshed model to the rig of the original model. Just assign skinning weight based on nearest neighbor :param filename_ori: original mesh filename :param filename_remesh: remeshed mesh filename :param pred_rig: predicted rig :return: predicted rig for original mesh """ mesh_remesh = o3d.io.read_triangle_mesh(filename_remesh) mesh_ori = o3d.io.read_triangle_mesh(filename_ori) tranfer_rig = Info() vert_remesh = np.asarray(mesh_remesh.vertices) vert_ori = np.asarray(mesh_ori.vertices) vertice_distance = np.sqrt( np.sum((vert_ori[np.newaxis, ...] - vert_remesh[:, np.newaxis, :])**2, axis=2)) vertice_raw_id = np.argmin( vertice_distance, axis=0 ) # nearest vertex id on the fixed mesh for each vertex on the remeshed mesh tranfer_rig.root = pred_rig.root tranfer_rig.joint_pos = pred_rig.joint_pos new_skin = [] for v in range(len(vert_ori)): skin_v = [v] v_nn = vertice_raw_id[v] skin_v += pred_rig.joint_skin[v_nn][1:] new_skin.append(skin_v) tranfer_rig.joint_skin = new_skin return tranfer_rig
# if i != args.character_idx: # continue for motion in motions_list[args.start:args.last]: joint_pos_file = os.path.join(joint_log, 'test', character, '%s_joint.npy' % (motion)) joint_pos = np.load(joint_pos_file, allow_pickle=True).item() joint_result = joint_pos[prediction_method] # save skeleton pred_skel = Info() nodes = [] for joint_index, joint_pos in enumerate(joint_result): nodes.append(TreeNode(name=joint_name[joint_index], pos=joint_pos)) pred_skel.root = nodes[0] for parent, children in enumerate(tree): for child in children: nodes[parent].children.append(nodes[child]) nodes[child].parent = nodes[parent] # calculate volumetric geodesic distance bones, _, _ = get_bones(pred_skel) mesh_filename = os.path.join(data_dir, 'objs', character + '/' + motion + '.obj') # mesh_filename = os.path.join(data_dir, 'test_objs', character + '_' + motion + '.obj') mesh = o3d.io.read_triangle_mesh(mesh_filename) mesh.compute_vertex_normals() mesh_v = np.asarray(mesh.vertices) surface_geodesic = calc_surface_geodesic(mesh) volumetric_geodesic = calc_geodesic_matrix(