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

vto conversion, some progress, not working properl

Browse Source
This commit is contained in:
luni
2023-05-29 10:28:31 +02:00
parent e34e2143e7
commit 3abe7554bc
1 changed files with 120 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
data/convert_vto.py
+120 -9
View File
@@ -6,13 +6,68 @@ 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))
@@ -62,10 +117,13 @@ def register_smpl():
# 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)
@@ -86,7 +144,10 @@ def register_smpl():
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)
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)
@@ -104,23 +165,56 @@ def register_smpl():
# 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)):
_, idx, _ = kd_tree.search_knn_vector_3d(aligned_source_vertices[i], 1)
correspondences.append((i, idx[0]))
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), ...]
# 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]]
# 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):
# files = os.listdir(folder)
files = ['/data-net/code/vto-dataset/tshirt/simulations/tshirt_shape00_55_27.pkl']
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:
@@ -131,10 +225,27 @@ def convert_vto(folder, out_folder):
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()
@@ -143,6 +254,6 @@ def convert_vto(folder, out_folder):
if __name__ == '__main__':
register_smpl()
# register_smpl()
convert_vto('/data-net/code/vto-dataset/dress/simulations', '/data-net/code/vto-dataset/dress_converted')
pass