131 lines
5.9 KiB
Python
131 lines
5.9 KiB
Python
import os
|
|
import sys
|
|
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
|
|
|
|
|
|
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = '0'
|
|
|
|
|
|
import time
|
|
import argparse
|
|
import random
|
|
import tqdm
|
|
import numpy as np
|
|
import plotly.graph_objects as go
|
|
import plotly.express as ex
|
|
from utils.postprocess import fix_collision
|
|
|
|
|
|
import torch
|
|
from tqdm import tqdm
|
|
from data.utils import Frame
|
|
from data.cloth3d_io import readPC2Frame, readPC2, readOBJ
|
|
from utils.utils import pickle_load, save_stl
|
|
from scipy.spatial import cKDTree
|
|
from smpl.smpl_torch import SMPL
|
|
from scipy.spatial.transform import Rotation as R
|
|
import trimesh as tm
|
|
from plotly.subplots import make_subplots
|
|
|
|
|
|
smpl = SMPL(path='smpl', device='cpu')
|
|
def deorient(verts, pose, inverse=True):
|
|
assert pose.shape[1] == 72
|
|
rodr = smpl.rodrigues(torch.tensor(pose, device=smpl.device).reshape(-1, 1, 3)).reshape(verts.shape[0], -1, 3, 3).numpy()
|
|
rotated = np.zeros_like(verts)
|
|
for frame in range(verts.shape[0]):
|
|
r = R.from_matrix(rodr[frame, 0, :, :])
|
|
rotated[frame] = r.apply(verts[frame], inverse=inverse)
|
|
|
|
return rotated
|
|
|
|
|
|
|
|
def calculate():
|
|
|
|
# garments = sorted(os.listdir(hood_folder))
|
|
# garments = [g[:-4] for g in garments]
|
|
# random.shuffle(garments)
|
|
garments = [('01909_Jumpsuit', 175), ('00960_Dress', 41), ('01573_Dress', 40), ('00079_Tshirt', 190)]
|
|
body_faces = np.load('smpl/smpl_faces.npy').astype(np.int32)
|
|
|
|
for garment_data in tqdm(garments):
|
|
garment = garment_data[0]
|
|
picker = 'skirt' if garment.split('_')[1] in ['Dress', 'Skirt'] else 'nonskirt'
|
|
id, garment_type = garment.split('_')
|
|
|
|
pred_hood_data = pickle_load(f'{hood_folder}/{garment}.pkl')
|
|
pred_hood, pred_hood_faces = pred_hood_data['verts'], pred_hood_data['faces']
|
|
gt = readPC2(os.path.join(cloth3d_folder, id, garment_type + '.pc16'), True)['V'][2:]
|
|
mask = np.load(f'/data-local/127_2/cloth3d_processed_test/mesh_fixed/{id}_{garment_type}.npz')['mask']
|
|
pose = np.load(f'/data-local/127_2/cloth3d_processed_test/pose/{id}.npy')[2:]
|
|
gt = deorient(gt, pose, True)
|
|
body = readPC2(f'/data-local/127_2/cloth3d_processed_test/mesh_body_posed_no_orient/{id}.pc16', True)['V'][2:]
|
|
fix_data = np.load(f'/data-local/127_2/cloth3d_processed_test/mesh_fixed/{id}_{garment_type}.npz')
|
|
faces = fix_data['faces']
|
|
|
|
|
|
# i = np.random.randint(0, len(pred_hood))
|
|
i = garment_data[1]
|
|
|
|
try:
|
|
if calculate_ours:
|
|
pred_ours = np.load(f'{ours_folder[picker]}/{id}_{i+2}_{garment_type}.npz', allow_pickle=True)['arr_0'].item()['lvl3']
|
|
except:
|
|
print('missed one')
|
|
continue
|
|
|
|
mesh_posed_fixed = fix_collision(torch.from_numpy(pred_ours[None,:,:]).to('cuda'), torch.from_numpy(body[i][None,:,:]).to('cuda'), torch.from_numpy(mask[None,:]).to('cuda')).cpu().numpy()
|
|
mesh_posed_fixed = mesh_posed_fixed[0, mask, :]
|
|
try:
|
|
m = tm.Trimesh(mesh_posed_fixed, faces, process=False, validate=False)
|
|
tm.smoothing.filter_humphrey(m, iterations=2)
|
|
pred_ours = m.vertices
|
|
except:
|
|
pred_ours = mesh_posed_fixed
|
|
print('failed to apply smoothing')
|
|
|
|
# fig = go.Figure(data=[go.Mesh3d(x=pred_ours[:,0], y=pred_ours[:,2], z=pred_ours[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='blue')])
|
|
# fig.add_trace(go.Mesh3d(x=body[i][:,0], y=body[i][:,2], z=body[i][:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray'))
|
|
# fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1], ), yaxis=dict(range=[-1, 1], ), zaxis=dict(range=[-1, 1], ), ), )
|
|
# fig.update_layout(scene_aspectmode='cube', scene_aspectratio=dict(x=1, y=1, z=1))
|
|
# fig.show()
|
|
|
|
fig = go.Figure(data=[go.Mesh3d(x=pred_hood[i][:,0], y=pred_hood[i][:,2], z=pred_hood[i][:,1], i=pred_hood_faces[:,0], j=pred_hood_faces[:,1], k=pred_hood_faces[:,2], color='blue')])
|
|
fig.add_trace(go.Mesh3d(x=body[i][:,0], y=body[i][:,2], z=body[i][:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray'))
|
|
fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1], ), yaxis=dict(range=[-1, 1], ), zaxis=dict(range=[-1, 1], ), ), )
|
|
fig.update_layout(scene_aspectmode='cube', scene_aspectratio=dict(x=1, y=1, z=1))
|
|
fig.show()
|
|
|
|
# fig = go.Figure(data=[go.Mesh3d(x=gt[i][:,0], y=gt[i][:,2], z=gt[i][:,1], i=pred_hood_faces[:,0], j=pred_hood_faces[:,1], k=pred_hood_faces[:,2], color='blue')])
|
|
# fig.add_trace(go.Mesh3d(x=body[i][:,0], y=body[i][:,2], z=body[i][:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2], color='gray'))
|
|
# fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1], ), yaxis=dict(range=[-1, 1], ), zaxis=dict(range=[-1, 1], ), ), )
|
|
# fig.update_layout(scene_aspectmode='cube', scene_aspectratio=dict(x=1, y=1, z=1))
|
|
# fig.show()
|
|
pass
|
|
|
|
# save_stl(f'temp/hood/2/{id}_{i}_{garment_type}_gt.stl', gt[i], pred_hood_faces)
|
|
# save_stl(f'temp/hood/2/{id}_{i}_{garment_type}_gt_body.stl', body[i], body_faces)
|
|
# save_stl(f'temp/hood/2/{id}_{i}_{garment_type}_hood.stl', pred_hood[i], pred_hood_faces)
|
|
# save_stl(f'temp/hood/ne2w/{id}_{i}_{garment_type}_hood_body.stl', body[i], body_faces)
|
|
# save_stl(f'temp/hood/2/{id}_{i}_{garment_type}_ours.stl', pred_ours, faces)
|
|
# save_stl(f'temp/hood/2/{id}_{i}_{garment_type}_ours_body.stl', body[i], body_faces)
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
|
hood_folder = '/data-local/127_2/results/hood/all_2_105k'
|
|
ours_folder = dict()
|
|
ours_folder['skirt'] = '/data-local/127_2/results/pyr_iccvw_64_128_256_512_skirt_e2e_200000'
|
|
ours_folder['nonskirt'] = '/data-local/127_2/results/pyr_iccvw_64_128_256_512_e2e_2_250000'
|
|
cloth3d_folder = '/data-net/dataset/cloth3d/test'
|
|
|
|
results_folder = '/data-local/127_2/results'
|
|
|
|
calculate_ours = True
|
|
|
|
|
|
|
|
calculate() |