diff --git a/data/convert_vto.py b/data/convert_vto.py new file mode 100644 index 0000000..85bd850 --- /dev/null +++ b/data/convert_vto.py @@ -0,0 +1,148 @@ +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 \ No newline at end of file diff --git a/data/dataset_vto.py b/data/dataset_vto.py new file mode 100644 index 0000000..09a3156 --- /dev/null +++ b/data/dataset_vto.py @@ -0,0 +1,40 @@ +import os +import sys + +import argparse +parser = argparse.ArgumentParser(description='') +parser.add_argument('--debug', action='store_true') +args = vars(parser.parse_args()) + +if args['debug']: + os.environ["CUDA_VISIBLE_DEVICES"] = '0' + +import math +import time +import torch +import random +import pytorch3d +import torchvision as tv +import numpy as np +import scipy.io as sio +import plotly.express as ex +from torch.utils.data import Dataset + +sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..') + +from data.utils import Frame + +class Dataset(Dataset): + def __init__(self, folder): + pass + + def __len__(self): + return len(self.frames) + + + def __getitem__(self, idx): + frame : Frame = self.frames[idx] + + +if __name__ == '__main__': + from torch.utils.data import DataLoader diff --git a/smpl/smpl_torch.py b/smpl/smpl_torch.py index fbac7fb..0be5575 100644 --- a/smpl/smpl_torch.py +++ b/smpl/smpl_torch.py @@ -206,9 +206,9 @@ class SMPL(Module): ) ) - joints = torch.matmul(results, torch.transpose(torch.concat([J, torch.ones((J.shape[0], 24, 1), dtype=torch.float32, device=J.get_device())], axis=-1)[:,:,None,:],2,3))[:,:,:3,0] + joints = torch.matmul(results, torch.transpose(torch.concat([J, torch.ones((J.shape[0], 24, 1), dtype=torch.float32, device=J.device)], axis=-1)[:,:,None,:],2,3))[:,:,:3,0] - T = torch.stack([torch.tensordot(results[i], self.weights[torch.greater(g, 0.5).to(torch.long)], dims=([0], [1])).permute(2, 0, 1) for i, g in enumerate(betas[:,-1])]) + T = torch.stack([torch.tensordot(results[i], self.weights_all[torch.greater(g, 0.5).to(torch.long)], dims=([0], [1])).permute(2, 0, 1) for i, g in enumerate(betas[:,-1])]) rest_shape_h = torch.cat( (v_posed, torch.ones((batch_num, v_posed.shape[1], 1), dtype=torch.float32).to(self.device)), dim=2 ) @@ -230,4 +230,4 @@ class SMPL(Module): if remove and verts is None: result = result[:,self.to_keep,:] - return result, joints, v_shaped, v_posed + return result, joints