55 lines
1.7 KiB
Python
55 lines
1.7 KiB
Python
import os
|
|
import numpy as np
|
|
import plotly.express as ex
|
|
from simpleicp import PointCloud, SimpleICP
|
|
from scipy.spatial.distance import cdist
|
|
|
|
folder = '/home/workspace/data/dataset/tailornet'
|
|
|
|
def get_data_split():
|
|
data = np.load(os.path.join(folder, 'split_static_pose_shape.npz'))
|
|
return data['train'], data['test']
|
|
|
|
|
|
def get_garment_info(garment_class):
|
|
data = np.load(os.path.join(folder, 'garment_class_info.pkl'), allow_pickle=True)
|
|
return data[garment_class]
|
|
|
|
|
|
def get_smpl():
|
|
data = np.load(os.path.join('TailorNet_dataset/smpl/smpl_female.npz'), allow_pickle=True)
|
|
data_high_res = np.load(os.path.join('TailorNet_dataset/smpl/smpl_hres_female.npz'), allow_pickle=True)
|
|
supersmpl = np.load(os.path.join('custom_utils/super_model_f.pkl'), allow_pickle=True)
|
|
|
|
return data_high_res['v_template'], supersmpl['v_template']
|
|
|
|
print()
|
|
|
|
|
|
def register_smpl(verts_1, verts_2):
|
|
|
|
# Create point cloud objects
|
|
pc_fix = PointCloud(verts_1, columns=["x", "y", "z"])
|
|
pc_mov = PointCloud(verts_2, columns=["x", "y", "z"])
|
|
|
|
# Create simpleICP object, add point clouds, and run algorithm!
|
|
icp = SimpleICP()
|
|
icp.add_point_clouds(pc_fix, pc_mov)
|
|
H, X_mov_transformed, rigid_body_transformation_params = icp.run(correspondences=10000, min_change=0.01, neighbors=25)
|
|
|
|
dist = cdist(verts_1, X_mov_transformed)
|
|
matching_dist = np.min(dist, axis=1)
|
|
matching = np.argmin(dist, axis=1)
|
|
|
|
|
|
print()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
train_split, test_split = get_data_split()
|
|
garment_info = get_garment_info('old-t-shirt')
|
|
smpl = get_smpl()
|
|
tailornet_smpl, super_smpl = get_smpl()
|
|
# registration = register_smpl(super_smpl, tailornet_smpl)
|
|
|
|
print() |