save all
This commit is contained in:
@@ -0,0 +1,148 @@
|
||||
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 readBIN(fname):
|
||||
if '.' not in fname: fname += '.bin'
|
||||
elif '.' in fname and not fname.endswith('.bin'):
|
||||
print("File name extension should be '.bin'")
|
||||
return
|
||||
|
||||
with open(fname, 'rb') as f:
|
||||
F = np.frombuffer(f.read(), dtype=np.uint16).astype(np.int32)
|
||||
return F.reshape((-1,3))
|
||||
|
||||
|
||||
def calculate():
|
||||
|
||||
garments = sorted(os.listdir('/data-local/127_2/cloth3d_processed_test/mesh_mask'))
|
||||
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 in tqdm(garments):
|
||||
# garment = garment_data[0]
|
||||
|
||||
picker = 'skirt' if garment.split('_')[1] in ['Dress', 'Skirt'] else 'nonskirt'
|
||||
id, garment_type = garment.split('_')
|
||||
|
||||
if not os.path.exists(f'temp/cloth3d/{id}/outfit.pc16'):
|
||||
continue
|
||||
if not os.path.exists(os.path.join(cloth3d_folder, id, garment_type + '.pc16')):
|
||||
continue
|
||||
|
||||
pred_cloth3d = readPC2(f'temp/cloth3d/{id}/outfit.pc16', True)['V']
|
||||
pred_cloth3d_faces = readBIN(f'temp/cloth3d/{id}/outfit.bin')
|
||||
gt = readPC2(os.path.join(cloth3d_folder, id, garment_type + '.pc16'), True)['V']
|
||||
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')
|
||||
gt = deorient(gt, pose, True)
|
||||
# gt = deorient(gt, pose, True)
|
||||
body = readPC2(f'/data-local/127_2/cloth3d_processed_test/mesh_body_posed_no_orient/{id}.pc16', True)['V']
|
||||
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_cloth3d))
|
||||
# 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_cloth3d[i][:,0], y=pred_cloth3d[i][:,1], z=pred_cloth3d[i][:,2], i=pred_cloth3d_faces[:,0], j=pred_cloth3d_faces[:,1], k=pred_cloth3d_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/{id}_{i}_{garment_type}_gt.stl', gt[i], pred_hood_faces)
|
||||
# save_stl(f'temp/hood/{id}_{i}_{garment_type}_gt_body.stl', body[i], body_faces)
|
||||
# save_stl(f'temp/hood/new/{id}_{i}_{garment_type}_hood.stl', pred_hood[i], pred_hood_faces)
|
||||
# save_stl(f'temp/hood/new/{id}_{i}_{garment_type}_hood_body.stl', body[i], body_faces)
|
||||
# save_stl(f'temp/hood/{id}_{i}_{garment_type}_ours.stl', pred_ours, faces)
|
||||
# save_stl(f'temp/hood/{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()
|
||||
Reference in New Issue
Block a user