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(
