259 lines
12 KiB
Python
259 lines
12 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
|
|
import trimesh
|
|
# from vedo import trimesh as trimesh_plot
|
|
from scipy.spatial import KDTree
|
|
from scipy.sparse import lil_matrix
|
|
from scipy.optimize import minimize
|
|
|
|
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
|
|
|
|
from smpl.smpl_torch import SMPL
|
|
|
|
# import mask faces
|
|
from custom_utils.custom_utils import _maskFaces
|
|
|
|
smpl = SMPL(path='smpl')
|
|
|
|
def energy_function(params, source_points, target_points, indices):
|
|
num_points = source_points.shape[0]
|
|
num_dims = source_points.shape[1]
|
|
|
|
# Reshape the parameters
|
|
weights = params[:num_points * num_dims].reshape((num_points, num_dims))
|
|
translation = params[num_points * num_dims:num_points * num_dims + num_dims]
|
|
control_points = params[num_points * num_dims + num_dims:].reshape((-1, num_dims))
|
|
|
|
# Compute the transformed source points
|
|
transformed_points = source_points + np.dot(weights, control_points.T)
|
|
|
|
# Compute the TPS energy
|
|
tps_energy = 0.5 * np.sum(np.square(target_points[indices] - transformed_points))
|
|
|
|
return tps_energy
|
|
|
|
def non_rigid(source_points, target_points):
|
|
tree = KDTree(target_points)
|
|
distances, indices = tree.query(source_points)
|
|
|
|
# Initialize the optimization parameters
|
|
num_points = source_points.shape[0]
|
|
num_dims = source_points.shape[1]
|
|
initial_params = np.zeros(num_points * num_dims + num_dims + num_points)
|
|
|
|
# Define the optimization bounds
|
|
bounds = [(None, None)] * (num_points * num_dims)
|
|
bounds += [(None, None)] * num_dims
|
|
bounds += [(None, None)] * (num_points * num_dims)
|
|
|
|
# Perform the optimization
|
|
result = minimize(energy_function, initial_params, method='L-BFGS-B',
|
|
args=(source_points, target_points, indices), bounds=bounds)
|
|
|
|
# Get the optimized parameters
|
|
optimized_params = result.x
|
|
|
|
# Reshape the parameters
|
|
optimized_weights = optimized_params[:num_points * num_dims].reshape((num_points, num_dims))
|
|
optimized_translation = optimized_params[num_points * num_dims:num_points * num_dims + num_dims]
|
|
optimized_control_points = optimized_params[num_points * num_dims + num_dims:].reshape((-1, num_dims))
|
|
|
|
transformed_points = source_points + np.dot(optimized_weights, optimized_control_points.T)
|
|
return transformed_points
|
|
|
|
|
|
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()
|
|
|
|
# non_rigid(garment, smpl_out.numpy())
|
|
|
|
# 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))
|
|
# flip source and target mesh
|
|
source_mesh, target_mesh = target_mesh, source_mesh
|
|
|
|
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.075,
|
|
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 = []
|
|
matched = set()
|
|
for i in range(len(aligned_source_vertices)):
|
|
k = 1
|
|
_, idx, _ = kd_tree.search_knn_vector_3d(aligned_source_vertices[i], k)
|
|
# while idx[-1] in matched:
|
|
# k += 1
|
|
# _, idx, _ = kd_tree.search_knn_vector_3d(aligned_source_vertices[i], k)
|
|
# matched.add(idx[-1])
|
|
correspondences.append((i, idx[-1]))
|
|
|
|
# correspondences is a list of tuples where each tuple contains the corresponding vertex indices
|
|
# For example, correspondences = [(0, 10), (1, 5), (2, 15), ...]
|
|
|
|
# Visualize the correspondences with plotly
|
|
# fig = go.Figure(data=[go.Scatter3d(x=aligned_source_vertices[:,0], y=aligned_source_vertices[:,1], z=aligned_source_vertices[:,2], mode='markers', marker=dict(color='red', size=1))])
|
|
# fig.add_trace(go.Scatter3d(x=target_vertices[:,0], y=target_vertices[:,1], z=target_vertices[:,2], mode='markers', marker=dict(color='blue', size=1)))
|
|
# for corr in correspondences:
|
|
# fig.add_trace(go.Scatter3d(x=[aligned_source_vertices[corr[0],0], target_vertices[corr[1],0]], y=[aligned_source_vertices[corr[0],1], target_vertices[corr[1],1]], z=[aligned_source_vertices[corr[0],2], target_vertices[corr[1],2]], mode='lines', line=dict(color='green', width=1)))
|
|
# fig.update_layout(scene=dict(xaxis=dict(range=[-1,1],), yaxis=dict(range=[-1,1],), zaxis=dict(range=[-1,1],),),)
|
|
# fig.show()
|
|
|
|
registered = np.zeros((14475, 3))
|
|
mask = np.zeros((14475), dtype=bool)
|
|
for i in range(len(correspondences)):
|
|
registered[correspondences[i][1]] = aligned_source_vertices[correspondences[i][0]]
|
|
mask[correspondences[i][1]] = True
|
|
|
|
# F = np.load(f'smpl/smpl_faces.npy').astype(np.int32)
|
|
|
|
# registered_masked, masked_faces = _maskFaces(registered, np.asarray(target_mesh.triangles), mask)
|
|
registered_masked, masked_faces = registered, np.asarray(target_mesh.triangles)
|
|
|
|
# plot aligned source mesh with the target topology
|
|
fig = go.Figure(data=[go.Mesh3d(x=registered_masked[:,0], y=registered_masked[:,1], z=registered_masked[:,2], i=masked_faces[:,0], j=masked_faces[:,1], k=masked_faces[:,2])])
|
|
fig.update_layout(scene=dict(xaxis=dict(range=[-1,1],), yaxis=dict(range=[-1,1],), zaxis=dict(range=[-1,1],),),)
|
|
fig.show()
|
|
|
|
return smpl
|
|
|
|
|
|
def convert_vto(folder, out_folder):
|
|
os.makedirs(out_folder, exist_ok=True)
|
|
os.makedirs(os.path.join(out_folder, 'mesh_posed_no_orient'), exist_ok=True)
|
|
os.makedirs(os.path.join(out_folder, 'mesh_posed_body_no_orient'), exist_ok=True)
|
|
os.makedirs(os.path.join(out_folder, 'mesh_mask'), exist_ok=True)
|
|
os.makedirs(os.path.join(out_folder, 'pose_shape'), exist_ok=True)
|
|
F = np.load(f'smpl/smpl_faces.npy').astype(np.int32)
|
|
|
|
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]
|
|
# pose[:, :3] = 0
|
|
id = id.replace('_', '-')
|
|
subject = id.split('-')[2]
|
|
sequence = '_'.join(id.split('-')[2:])
|
|
# cmu_data = np.load(f'/data-net/code/vto-garment-collisions/assets/CMU/{subject}/{sequence}_poses.npz')
|
|
# shape = np.tile(cmu_data['betas'], (shape.shape[0], 1))
|
|
# pose = cmu_data['poses']
|
|
|
|
body = smpl(torch.tensor(pose), torch.tensor(np.concatenate((shape, np.zeros((shape.shape[0],1),dtype=np.float32)),axis=1)))[0].numpy()
|
|
pose_shape = np.concatenate((pose, shape, np.zeros((shape.shape[0],1),dtype=np.float32)),axis=1)
|
|
|
|
# for i in range(vertices.shape[0]):
|
|
# np.savez_compressed(os.path.join(out_folder, 'mesh_posed_body_no_orient', f'{id}_{i}.npz'), mesh_posed_body_no_orient=body[i])
|
|
# np.save(os.path.join(out_folder, 'pose_shape', f'{id}_{i}.npy'), pose_shape[i])
|
|
|
|
|
|
# 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.add_trace(go.Mesh3d(x=body[frame,:,0], y=body[frame,:,1], z=body[frame,:,2], i=F[:,0], j=F[:,1], k=F[:,2]))
|
|
fig.update_layout(scene=dict(xaxis=dict(range=[-1,1],), yaxis=dict(range=[-1,1],), zaxis=dict(range=[-1,1],),),)
|
|
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 |