From 3abe7554bc5524142d105a0976c2c923732005af Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Mon, 29 May 2023 10:28:31 +0200 Subject: [PATCH] vto conversion, some progress, not working properl --- data/convert_vto.py | 129 ++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 120 insertions(+), 9 deletions(-) diff --git a/data/convert_vto.py b/data/convert_vto.py index 85bd850..e998e9b 100644 --- a/data/convert_vto.py +++ b/data/convert_vto.py @@ -6,13 +6,68 @@ import plotly.express as ex import plotly.graph_objects as go import torch import open3d as o3d +import trimesh +# from vedo import trimesh as trimesh_plot +from scipy.spatial import KDTree +from scipy.sparse import lil_matrix +from scipy.optimize import minimize sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..') from smpl.smpl_torch import SMPL +# import mask faces +from custom_utils.custom_utils import _maskFaces + smpl = SMPL(path='smpl') +def energy_function(params, source_points, target_points, indices): + num_points = source_points.shape[0] + num_dims = source_points.shape[1] + + # Reshape the parameters + weights = params[:num_points * num_dims].reshape((num_points, num_dims)) + translation = params[num_points * num_dims:num_points * num_dims + num_dims] + control_points = params[num_points * num_dims + num_dims:].reshape((-1, num_dims)) + + # Compute the transformed source points + transformed_points = source_points + np.dot(weights, control_points.T) + + # Compute the TPS energy + tps_energy = 0.5 * np.sum(np.square(target_points[indices] - transformed_points)) + + return tps_energy + +def non_rigid(source_points, target_points): + tree = KDTree(target_points) + distances, indices = tree.query(source_points) + + # Initialize the optimization parameters + num_points = source_points.shape[0] + num_dims = source_points.shape[1] + initial_params = np.zeros(num_points * num_dims + num_dims + num_points) + + # Define the optimization bounds + bounds = [(None, None)] * (num_points * num_dims) + bounds += [(None, None)] * num_dims + bounds += [(None, None)] * (num_points * num_dims) + + # Perform the optimization + result = minimize(energy_function, initial_params, method='L-BFGS-B', + args=(source_points, target_points, indices), bounds=bounds) + + # Get the optimized parameters + optimized_params = result.x + + # Reshape the parameters + optimized_weights = optimized_params[:num_points * num_dims].reshape((num_points, num_dims)) + optimized_translation = optimized_params[num_points * num_dims:num_points * num_dims + num_dims] + optimized_control_points = optimized_params[num_points * num_dims + num_dims:].reshape((-1, num_dims)) + + transformed_points = source_points + np.dot(optimized_weights, optimized_control_points.T) + return transformed_points + + def register_smpl(): # generate body with zero pose and shape pose = torch.zeros((1, 72)) @@ -62,10 +117,13 @@ def register_smpl(): # fig.update_layout(scene=dict(xaxis=dict(range=[-1,1],), yaxis=dict(range=[-1,1],), zaxis=dict(range=[-1,1],),),) # fig.show() + # non_rigid(garment, smpl_out.numpy()) # Load the source and target meshes source_mesh = o3d.geometry.TriangleMesh(o3d.utility.Vector3dVector(garment), o3d.utility.Vector3iVector(faces)) target_mesh = o3d.geometry.TriangleMesh(o3d.utility.Vector3dVector(smpl_out), o3d.utility.Vector3iVector(smpl.f)) + # flip source and target mesh + source_mesh, target_mesh = target_mesh, source_mesh source_mesh.compute_vertex_normals() source_mesh = source_mesh.filter_smooth_laplacian(number_of_iterations=10) @@ -86,7 +144,10 @@ def register_smpl(): criteria = o3d.pipelines.registration.ICPConvergenceCriteria(max_iteration=2000) # Perform ICP registration - reg_result = o3d.pipelines.registration.registration_icp(o3d.geometry.PointCloud(source_mesh.vertices), o3d.geometry.PointCloud(target_mesh.vertices), 0.01, init_transformation, icp, criteria) + reg_result = o3d.pipelines.registration.registration_icp(o3d.geometry.PointCloud(source_mesh.vertices), + o3d.geometry.PointCloud(target_mesh.vertices), + 0.075, + init_transformation, icp, criteria) # Get the aligned source mesh after registration aligned_mesh = source_mesh.transform(reg_result.transformation) @@ -104,23 +165,56 @@ def register_smpl(): # Perform nearest neighbor search to find corresponding vertex indices kd_tree = o3d.geometry.KDTreeFlann(target_mesh) correspondences = [] + matched = set() for i in range(len(aligned_source_vertices)): - _, idx, _ = kd_tree.search_knn_vector_3d(aligned_source_vertices[i], 1) - correspondences.append((i, idx[0])) + k = 1 + _, idx, _ = kd_tree.search_knn_vector_3d(aligned_source_vertices[i], k) + # while idx[-1] in matched: + # k += 1 + # _, idx, _ = kd_tree.search_knn_vector_3d(aligned_source_vertices[i], k) + # matched.add(idx[-1]) + correspondences.append((i, idx[-1])) # correspondences is a list of tuples where each tuple contains the corresponding vertex indices # For example, correspondences = [(0, 10), (1, 5), (2, 15), ...] - # You can access the corresponding vertices using the vertex indices - corresponding_source_vertices = aligned_source_vertices[correspondences[:, 0]] - corresponding_target_vertices = target_vertices[correspondences[:, 1]] + # Visualize the correspondences with plotly + # fig = go.Figure(data=[go.Scatter3d(x=aligned_source_vertices[:,0], y=aligned_source_vertices[:,1], z=aligned_source_vertices[:,2], mode='markers', marker=dict(color='red', size=1))]) + # fig.add_trace(go.Scatter3d(x=target_vertices[:,0], y=target_vertices[:,1], z=target_vertices[:,2], mode='markers', marker=dict(color='blue', size=1))) + # for corr in correspondences: + # fig.add_trace(go.Scatter3d(x=[aligned_source_vertices[corr[0],0], target_vertices[corr[1],0]], y=[aligned_source_vertices[corr[0],1], target_vertices[corr[1],1]], z=[aligned_source_vertices[corr[0],2], target_vertices[corr[1],2]], mode='lines', line=dict(color='green', width=1))) + # fig.update_layout(scene=dict(xaxis=dict(range=[-1,1],), yaxis=dict(range=[-1,1],), zaxis=dict(range=[-1,1],),),) + # fig.show() + + registered = np.zeros((14475, 3)) + mask = np.zeros((14475), dtype=bool) + for i in range(len(correspondences)): + registered[correspondences[i][1]] = aligned_source_vertices[correspondences[i][0]] + mask[correspondences[i][1]] = True + + # F = np.load(f'smpl/smpl_faces.npy').astype(np.int32) + + # registered_masked, masked_faces = _maskFaces(registered, np.asarray(target_mesh.triangles), mask) + registered_masked, masked_faces = registered, np.asarray(target_mesh.triangles) + + # plot aligned source mesh with the target topology + fig = go.Figure(data=[go.Mesh3d(x=registered_masked[:,0], y=registered_masked[:,1], z=registered_masked[:,2], i=masked_faces[:,0], j=masked_faces[:,1], k=masked_faces[:,2])]) + fig.update_layout(scene=dict(xaxis=dict(range=[-1,1],), yaxis=dict(range=[-1,1],), zaxis=dict(range=[-1,1],),),) + fig.show() return smpl def convert_vto(folder, out_folder): - # files = os.listdir(folder) - files = ['/data-net/code/vto-dataset/tshirt/simulations/tshirt_shape00_55_27.pkl'] + os.makedirs(out_folder, exist_ok=True) + os.makedirs(os.path.join(out_folder, 'mesh_posed_no_orient'), exist_ok=True) + os.makedirs(os.path.join(out_folder, 'mesh_posed_body_no_orient'), exist_ok=True) + os.makedirs(os.path.join(out_folder, 'mesh_mask'), exist_ok=True) + os.makedirs(os.path.join(out_folder, 'pose_shape'), exist_ok=True) + F = np.load(f'smpl/smpl_faces.npy').astype(np.int32) + + files = os.listdir(folder) + # files = ['/data-net/code/vto-dataset/tshirt/simulations/tshirt_shape00_55_27.pkl'] for file in files: # read pickle files and get vetices and faces from them with open(os.path.join(folder, file), "rb") as f: @@ -131,10 +225,27 @@ def convert_vto(folder, out_folder): shape = data['shape'] trans = data['translation'] id = file.split('.')[0] + # pose[:, :3] = 0 + id = id.replace('_', '-') + subject = id.split('-')[2] + sequence = '_'.join(id.split('-')[2:]) + # cmu_data = np.load(f'/data-net/code/vto-garment-collisions/assets/CMU/{subject}/{sequence}_poses.npz') + # shape = np.tile(cmu_data['betas'], (shape.shape[0], 1)) + # pose = cmu_data['poses'] + + body = smpl(torch.tensor(pose), torch.tensor(np.concatenate((shape, np.zeros((shape.shape[0],1),dtype=np.float32)),axis=1)))[0].numpy() + pose_shape = np.concatenate((pose, shape, np.zeros((shape.shape[0],1),dtype=np.float32)),axis=1) + + # for i in range(vertices.shape[0]): + # np.savez_compressed(os.path.join(out_folder, 'mesh_posed_body_no_orient', f'{id}_{i}.npz'), mesh_posed_body_no_orient=body[i]) + # np.save(os.path.join(out_folder, 'pose_shape', f'{id}_{i}.npy'), pose_shape[i]) + # plot mesh with plotly frame = 0 fig = go.Figure(data=[go.Mesh3d(x=vertices[frame,:,0], y=vertices[frame,:,1], z=vertices[frame,:,2], i=faces[:,0], j=faces[:,1], k=faces[:,2])]) + fig.add_trace(go.Mesh3d(x=body[frame,:,0], y=body[frame,:,1], z=body[frame,:,2], i=F[:,0], j=F[:,1], k=F[:,2])) + fig.update_layout(scene=dict(xaxis=dict(range=[-1,1],), yaxis=dict(range=[-1,1],), zaxis=dict(range=[-1,1],),),) fig.show() @@ -143,6 +254,6 @@ def convert_vto(folder, out_folder): if __name__ == '__main__': - register_smpl() + # register_smpl() convert_vto('/data-net/code/vto-dataset/dress/simulations', '/data-net/code/vto-dataset/dress_converted') pass \ No newline at end of file