initial work
This commit is contained in:
@@ -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
|
||||
@@ -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
|
||||
+3
-3
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user