initial work

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

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

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