def genDataset(process_id):
global dataset_folder
print("process ID {:d}".format(process_id))
if process_id < 6:
model_list = np.loadtxt(os.path.join(dataset_folder,
'train_final.txt'),
dtype=int)
model_list = model_list[365 * process_id:365 * (process_id + 1)]
split_name = 'train'
elif process_id == 6:
model_list = np.loadtxt(os.path.join(dataset_folder, 'val_final.txt'),
dtype=int)
split_name = 'val'
elif process_id == 7:
model_list = np.loadtxt(os.path.join(dataset_folder, 'test_final.txt'),
dtype=int)
split_name = 'test'
mkdir_p(os.path.join(dataset_folder, split_name))
for model_id in model_list:
remeshed_obj_filename = os.path.join(
dataset_folder, 'obj_remesh/{:d}.obj'.format(model_id))
info_filename = os.path.join(
dataset_folder, 'rig_info_remesh/{:d}.txt'.format(model_id))
remeshed_obj = o3d.io.read_triangle_mesh(remeshed_obj_filename)
remesh_obj_v = np.asarray(remeshed_obj.vertices)
remesh_obj_vn = np.asarray(remeshed_obj.vertex_normals)
remesh_obj_f = np.asarray(remeshed_obj.triangles)
rig_info = Info(info_filename)
#vertices
vert_filename = os.path.join(
dataset_folder, '{:s}/{:d}_v.txt'.format(split_name, model_id))
input_feature = np.concatenate((remesh_obj_v, remesh_obj_vn), axis=1)
np.savetxt(vert_filename, input_feature, fmt='%.6f')
#topology edges
edge_index = get_tpl_edges(remesh_obj_v, remesh_obj_f)
graph_filename = os.path.join(
dataset_folder, '{:s}/{:d}_tpl_e.txt'.format(split_name, model_id))
np.savetxt(graph_filename, edge_index, fmt='%d')
# geodesic_edges
surface_geodesic = calc_surface_geodesic(remeshed_obj)
edge_index = get_geo_edges(surface_geodesic, remesh_obj_v)
graph_filename = os.path.join(
dataset_folder, '{:s}/{:d}_geo_e.txt'.format(split_name, model_id))
np.savetxt(graph_filename, edge_index, fmt='%d')
# joints
joint_pos = rig_info.get_joint_dict()
joint_name_list = list(joint_pos.keys())
joint_pos_list = list(joint_pos.values())
joint_pos_list = [np.array(i) for i in joint_pos_list]
adjacent_matrix = rig_info.adjacent_matrix()
joint_filename = os.path.join(
dataset_folder, '{:s}/{:d}_j.txt'.format(split_name, model_id))
adj_filename = os.path.join(
dataset_folder, '{:s}/{:d}_adj.txt'.format(split_name, model_id))
np.savetxt(adj_filename, adjacent_matrix, fmt='%d')
np.savetxt(joint_filename, np.array(joint_pos_list), fmt='%.6f')
# pre_trained attn
shutil.copyfile(
os.path.join(dataset_folder,
'pretrain_attention/{:d}.txt'.format(model_id)),
os.path.join(dataset_folder,
'{:s}/{:d}_attn.txt'.format(split_name, model_id)))
# voxel
shutil.copyfile(
os.path.join(dataset_folder, 'vox/{:d}.binvox'.format(model_id)),
os.path.join(dataset_folder,
'{:s}/{:d}.binvox'.format(split_name, model_id)))
#skinning information
num_nearest_bone = 5
geo_dist = np.load(
os.path.join(
dataset_folder,
"volumetric_geodesic/{:d}_volumetric_geo.npy".format(
model_id)))
bone_pos, bone_names, bone_isleaf = get_bones(rig_info)
input_samples = [] # mesh_vertex_id, (bone_id, 1 / D_g, is_leaf) * N
ground_truth_labels = [] # w_1, w_2, ..., w_N
for vert_remesh_id in range(len(remesh_obj_v)):
this_sample = [vert_remesh_id]
this_label = []
skin = rig_info.joint_skin[vert_remesh_id]
skin_w = {}
for i in np.arange(1, len(skin), 2):
skin_w[skin[i]] = float(skin[i + 1])
bone_id_near_to_far = np.argsort(geo_dist[vert_remesh_id, :])
for i in range(num_nearest_bone):
if i >= len(bone_id_near_to_far):
this_sample += [-1, 0, 0]
this_label.append(0.0)
continue
bone_id = bone_id_near_to_far[i]
this_sample.append(bone_id)
this_sample.append(1.0 /
(geo_dist[vert_remesh_id, bone_id] + 1e-10))
this_sample.append(bone_isleaf[bone_id])
start_joint_name = bone_names[bone_id][0]
if start_joint_name in skin_w:
this_label.append(skin_w[start_joint_name])
del skin_w[start_joint_name]
else:
this_label.append(0.0)
input_samples.append(this_sample)
ground_truth_labels.append(this_label)
with open(
os.path.join(dataset_folder, '{:s}/{:d}_skin.txt').format(
split_name, model_id), 'w') as fout:
for i in range(len(bone_pos)):
fout.write('bones {:s} {:s} {:.6f} {:.6f} {:.6f} '
'{:.6f} {:.6f} {:.6f}\n'.format(
bone_names[i][0], bone_names[i][1],
bone_pos[i, 0], bone_pos[i, 1], bone_pos[i, 2],
bone_pos[i, 3], bone_pos[i, 4], bone_pos[i, 5]))
for i in range(len(input_samples)):
fout.write('bind {:d} '.format(input_samples[i][0]))
for j in np.arange(1, len(input_samples[i]), 3):
fout.write('{:d} {:.6f} {:d} '.format(
input_samples[i][j], input_samples[i][j + 1],
input_samples[i][j + 2]))
fout.write('\n')
for i in range(len(ground_truth_labels)):
fout.write('influence ')
for j in range(len(ground_truth_labels[i])):
fout.write('{:.3f} '.format(ground_truth_labels[i][j]))
fout.write('\n')
def predict_skinning(input_data,
pred_skel,
skin_pred_net,
surface_geodesic,
mesh_filename,
subsampling=False):
"""
predict skinning
:param input_data: wrapped input data
:param pred_skel: predicted skeleton
:param skin_pred_net: network to predict skinning weights
:param surface_geodesic: geodesic distance matrix of all vertices
:param mesh_filename: mesh filename
:return: predicted rig with skinning weights information
"""
global device, output_folder
num_nearest_bone = 5
bones, bone_names, bone_isleaf = get_bones(pred_skel)
mesh_v = input_data.pos.data.cpu().numpy()
print(
" calculating volumetric geodesic distance from vertices to bone. This step takes some time..."
)
geo_dist = calc_geodesic_matrix(bones,
mesh_v,
surface_geodesic,
mesh_filename,
subsampling=subsampling)
input_samples = [] # joint_pos (x, y, z), (bone_id, 1/D)*5
loss_mask = []
skin_nn = []
for v_id in range(len(mesh_v)):
geo_dist_v = geo_dist[v_id]
bone_id_near_to_far = np.argsort(geo_dist_v)
this_sample = []
this_nn = []
this_mask = []
for i in range(num_nearest_bone):
if i >= len(bones):
this_sample += bones[bone_id_near_to_far[0]].tolist()
this_sample.append(
1.0 / (geo_dist_v[bone_id_near_to_far[0]] + 1e-10))
this_sample.append(bone_isleaf[bone_id_near_to_far[0]])
this_nn.append(0)
this_mask.append(0)
else:
skel_bone_id = bone_id_near_to_far[i]
this_sample += bones[skel_bone_id].tolist()
this_sample.append(1.0 / (geo_dist_v[skel_bone_id] + 1e-10))
this_sample.append(bone_isleaf[skel_bone_id])
this_nn.append(skel_bone_id)
this_mask.append(1)
input_samples.append(np.array(this_sample)[np.newaxis, :])
skin_nn.append(np.array(this_nn)[np.newaxis, :])
loss_mask.append(np.array(this_mask)[np.newaxis, :])
skin_input = np.concatenate(input_samples, axis=0)
loss_mask = np.concatenate(loss_mask, axis=0)
skin_nn = np.concatenate(skin_nn, axis=0)
skin_input = torch.from_numpy(skin_input).float()
input_data.skin_input = skin_input
input_data.to(device)
skin_pred = skin_pred_net(data)
skin_pred = torch.softmax(skin_pred, dim=1)
skin_pred = skin_pred.data.cpu().numpy()
skin_pred = skin_pred * loss_mask
skin_nn = skin_nn[:, 0:num_nearest_bone]
skin_pred_full = np.zeros((len(skin_pred), len(bone_names)))
for v in range(len(skin_pred)):
for nn_id in range(len(skin_nn[v, :])):
skin_pred_full[v, skin_nn[v, nn_id]] = skin_pred[v, nn_id]
print(" filtering skinning prediction")
tpl_e = input_data.tpl_edge_index.data.cpu().numpy()
skin_pred_full = post_filter(skin_pred_full, tpl_e, num_ring=1)
skin_pred_full[skin_pred_full <
np.max(skin_pred_full, axis=1, keepdims=True) * 0.35] = 0.0
skin_pred_full = skin_pred_full / (
skin_pred_full.sum(axis=1, keepdims=True) + 1e-10)
skel_res = assemble_skel_skin(pred_skel, skin_pred_full)
return skel_res
def one_process(dataset_folder, start_id, end_id):
model_list = np.loadtxt(os.path.join(dataset_folder, 'model_list.txt'),
dtype=int)
model_list = model_list[start_id:end_id]
remesh_obj_folder = os.path.join(dataset_folder, "obj_remesh")
mkdir_p(os.path.join(dataset_folder, "volumetric_geodesic/"))
for model_id in model_list:
print(model_id)
if os.path.exists(
os.path.join(
dataset_folder,
"volumetric_geodesic/{:d}_volumetric_geo.npy".format(
model_id))):
continue
remeshed_obj_filename = os.path.join(
dataset_folder, 'obj_remesh/{:d}.obj'.format(model_id))
ori_obj_filename = os.path.join(dataset_folder,
'obj/{:d}.obj'.format(model_id))
info_filename = os.path.join(dataset_folder,
'rig_info/{:d}.txt'.format(model_id))
pts = np.array(
o3d.io.read_triangle_mesh(
os.path.join(remesh_obj_folder,
'{:d}.obj'.format(model_id))).vertices)
mesh_remesh = trimesh.load(remeshed_obj_filename)
mesh_ori = trimesh.load(ori_obj_filename)
rig_info = Info(info_filename)
bones, bone_name, _ = get_bones(rig_info)
origins, ends, pts_bone_dist = pts2line(pts, bones)
if os.path.exists(
os.path.join(
dataset_folder,
"volumetric_geodesic/{:d}_visibility_raw.npy".format(
model_id))):
pts_bone_visibility = np.load(
os.path.join(
dataset_folder,
"volumetric_geodesic/{:d}_visibility_raw.npy".format(
model_id)))
else:
# pick one mesh with fewer faces to speed up
if len(mesh_remesh.faces) < len(mesh_ori.faces):
trimesh.repair.fix_normals(mesh_remesh)
pts_bone_visibility = calc_pts2bone_visible_mat(
mesh_remesh, origins, ends)
else:
trimesh.repair.fix_normals(mesh_ori)
pts_bone_visibility = calc_pts2bone_visible_mat(
mesh_ori, origins, ends)
pts_bone_visibility = pts_bone_visibility.reshape(
len(bones), len(pts)).transpose()
#np.save(os.path.join(dataset_folder, "volumetric_geodesic/{:d}_visibility_raw.npy".format(model_id)), pts_bone_visibility)
pts_bone_dist = pts_bone_dist.reshape(len(bones), len(pts)).transpose()
# remove visible points which are too far
if os.path.exists(
os.path.join(
dataset_folder,
"volumetric_geodesic/{:d}_visibility_filtered.npy".format(
model_id))):
pts_bone_visibility = np.load(
os.path.join(
dataset_folder,
"volumetric_geodesic/{:d}_visibility_filtered.npy".format(
model_id)))
else:
for b in range(pts_bone_visibility.shape[1]):
visible_pts = np.argwhere(
pts_bone_visibility[:, b] == 1).squeeze(1)
if len(visible_pts) == 0:
continue
threshold_b = np.percentile(pts_bone_dist[visible_pts, b], 15)
pts_bone_visibility[pts_bone_dist[:, b] > 1.3 * threshold_b,
b] = False
#np.save(os.path.join(dataset_folder, "volumetric_geodesic/{:d}_visibility_filtered.npy".format(model_id)), pts_bone_visibility)
mesh = o3d.io.read_triangle_mesh(
os.path.join(remesh_obj_folder, '{:d}.obj'.format(model_id)))
surface_geodesic = calc_surface_geodesic(mesh)
visible_matrix = np.zeros(pts_bone_visibility.shape)
visible_matrix[np.where(
pts_bone_visibility == 1)] = pts_bone_dist[np.where(
pts_bone_visibility == 1)]
euc_dist = np.sqrt(
np.sum((pts[np.newaxis, ...] - pts[:, np.newaxis, :])**2, axis=2))
for c in range(visible_matrix.shape[1]):
unvisible_pts = np.argwhere(pts_bone_visibility[:,
c] == 0).squeeze(1)
visible_pts = np.argwhere(pts_bone_visibility[:,
c] == 1).squeeze(1)
if len(visible_pts) == 0:
visible_matrix[:, c] = pts_bone_dist[:, c]
continue
for r in unvisible_pts:
dist1 = np.min(surface_geodesic[r, visible_pts])
nn_visible = visible_pts[np.argmin(
surface_geodesic[r, visible_pts])]
if np.isinf(dist1):
visible_matrix[r, c] = 8.0 + pts_bone_dist[r, c]
else:
visible_matrix[r,
c] = dist1 + visible_matrix[nn_visible, c]
np.save(
os.path.join(
dataset_folder,
"volumetric_geodesic/{:d}_volumetric_geo.npy".format(
model_id)), visible_matrix)
