import numpy as np import os import sys import pickle import plotly.express as ex import plotly.graph_objects as go import torch import open3d as o3d sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..') from smpl.smpl_torch import SMPL smpl = SMPL(path='smpl') def register_smpl(): # generate body with zero pose and shape pose = torch.zeros((1, 72)) shape = torch.zeros((1, 11)) smpl_out = smpl(pose, shape)[0][0] with open('/data-net/code/vto-dataset/tshirt/simulations/tshirt_shape00_55_27.pkl', "rb") as f: data = pickle.load(f) garment = data['vertices'][100] faces = data['faces'] # rotate garment x degrees along y axis # convert degrees to radians degrees = 107 theta = degrees * np.pi / 180 rot = np.array([[np.cos(theta), 0, np.sin(theta)], [0, 1, 0], [-np.sin(theta), 0, np.cos(theta)]]) garment = np.dot(rot, garment.T).T # rotate garment x degrees along x axis # convert degrees to radians degrees = 5 theta = degrees * np.pi / 180 rot = np.array([[1, 0, 0], [0, np.cos(theta), -np.sin(theta)], [0, np.sin(theta), np.cos(theta)]]) garment = np.dot(rot, garment.T).T # rotate garment x degrees along z axis # convert degrees to radians degrees = -5 theta = degrees * np.pi / 180 rot = np.array([[np.cos(theta), -np.sin(theta), 0], [np.sin(theta), np.cos(theta), 0], [0, 0, 1]]) garment = np.dot(rot, garment.T).T garment[:,0] -= 0.06 garment[:,1] -= 0.03 garment[:,2] -= 0.015 # visualize the body and garment # fig = go.Figure(data=[go.Mesh3d(x=smpl_out[:,0], y=smpl_out[:,1], z=smpl_out[:,2], i=smpl.f[:,0], j=smpl.f[:,1], k=smpl.f[:,2])]) # fig.add_trace(go.Mesh3d(x=garment[:,0], y=garment[:,1], z=garment[:,2], i=faces[:,0], j=faces[:,1], k=faces[:,2])) # fig.update_layout(scene=dict(xaxis=dict(range=[-1,1],), yaxis=dict(range=[-1,1],), zaxis=dict(range=[-1,1],),),) # fig.show() # 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)) source_mesh.compute_vertex_normals() source_mesh = source_mesh.filter_smooth_laplacian(number_of_iterations=10) # plot the smoothed mesh with plotly # fig = go.Figure(data=[go.Mesh3d(x=np.asarray(source_mesh.vertices)[:,0], y=np.asarray(source_mesh.vertices)[:,1], z=np.asarray(source_mesh.vertices)[:,2], i=np.asarray(source_mesh.triangles)[:,0], j=np.asarray(source_mesh.triangles)[:,1], k=np.asarray(source_mesh.triangles)[:,2])]) # fig.show() # Perform initial alignment (translation and/or scaling) # You can modify this step based on your specific needs # For example, you can manually set an initial transformation init_transformation = [[1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0], [0.0, 0.0, 0.0, 1.0]] # Create the ICP object and set the parameters icp = o3d.pipelines.registration.TransformationEstimationPointToPoint() 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) # Get the aligned source mesh after registration aligned_mesh = source_mesh.transform(reg_result.transformation) # Visualize the aligned mesh and target mesh with plotly fig = go.Figure(data=[go.Mesh3d(x=np.asarray(aligned_mesh.vertices)[:,0], y=np.asarray(aligned_mesh.vertices)[:,1], z=np.asarray(aligned_mesh.vertices)[:,2], i=np.asarray(aligned_mesh.triangles)[:,0], j=np.asarray(aligned_mesh.triangles)[:,1], k=np.asarray(aligned_mesh.triangles)[:,2])]) fig.add_trace(go.Mesh3d(x=np.asarray(target_mesh.vertices)[:,0], y=np.asarray(target_mesh.vertices)[:,1], z=np.asarray(target_mesh.vertices)[:,2], i=np.asarray(target_mesh.triangles)[:,0], j=np.asarray(target_mesh.triangles)[:,1], k=np.asarray(target_mesh.triangles)[:,2])) fig.update_layout(scene=dict(xaxis=dict(range=[-1,1],), yaxis=dict(range=[-1,1],), zaxis=dict(range=[-1,1],),),) fig.show() # Get the vertices of the aligned source mesh and target mesh aligned_source_vertices = np.asarray(aligned_mesh.vertices) target_vertices = np.asarray(target_mesh.vertices) # Perform nearest neighbor search to find corresponding vertex indices kd_tree = o3d.geometry.KDTreeFlann(target_mesh) correspondences = [] 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])) # 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]] 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'] for file in files: # read pickle files and get vetices and faces from them with open(os.path.join(folder, file), "rb") as f: data = pickle.load(f) vertices = data['vertices'] faces = data['faces'] pose = data['pose'] shape = data['shape'] trans = data['translation'] id = file.split('.')[0] # 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.show() pass pass if __name__ == '__main__': register_smpl() convert_vto('/data-net/code/vto-dataset/dress/simulations', '/data-net/code/vto-dataset/dress_converted') pass