148 lines
6.3 KiB
Python
148 lines
6.3 KiB
Python
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 |