luni/cvpr_2023
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

initial work

Browse Source
This commit is contained in:
luni
2023-05-25 10:57:55 +02:00
parent 53217c6e5d
commit e34e2143e7
3 changed files with 191 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
data/convert_vto.py
+148
View File
@@ -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