263 lines
13 KiB
Python
263 lines
13 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"] = '6'
|
|
|
|
|
|
import time
|
|
import argparse
|
|
|
|
import tqdm
|
|
import numpy as np
|
|
import plotly.graph_objects as go
|
|
import plotly.express as ex
|
|
|
|
|
|
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
|
|
from scipy.spatial import cKDTree
|
|
from smpl.smpl_torch import SMPL
|
|
from scipy.spatial.transform import Rotation as R
|
|
import trimesh as tm
|
|
from losses.losses import collision_loss, normal_loss, kl_loss, collision_error
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("-s", "--start", help = "start of range", default=0, type=int)
|
|
parser.add_argument("-e", "--end", help = "end of range", default=550, type=int)
|
|
parser.add_argument("-n", "--name", help = "end of range", default='lhunor')
|
|
parser.add_argument('--debug', action='store_true')
|
|
args = parser.parse_args()
|
|
|
|
def chamfer(points1, points2, points2_tree: cKDTree):
|
|
tree1 = cKDTree(points1)
|
|
dist1, _ = tree1.query(points2, workers=10)
|
|
dist2, _ = points2_tree.query(points1, workers=10)
|
|
chamfer_dist = np.mean(dist1) + np.mean(dist2)
|
|
return chamfer_dist
|
|
|
|
|
|
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 get_collision_depth(m1, m2):
|
|
cm = tm.collision.CollisionManager()
|
|
cm.add_object('m1', m1)
|
|
cm.add_object('m2', m2)
|
|
data = cm.in_collision_internal(return_data=True)[1]
|
|
# c = np.asarray([d.index('m2') for d in data])
|
|
# count = np.unique(c).shape[0]
|
|
depth_sum = np.sum([d.depth for d in data])
|
|
return depth_sum
|
|
|
|
|
|
def get_area_diff(m1, m2):
|
|
return np.abs(m1.area - m2.area)
|
|
|
|
device = 'cuda'
|
|
def get_collision_loss(pred, body, mask):
|
|
c_loss, c_ratio = collision_error(torch.from_numpy(pred).unsqueeze(0).to(device), torch.from_numpy(body).unsqueeze(0).to(device), torch.from_numpy(mask).unsqueeze(0).to(device))
|
|
return c_loss.cpu().numpy()[0], c_ratio.cpu().numpy()
|
|
|
|
def get_collision_loss_2(pred, body, mask):
|
|
_pred = np.zeros((14475, 3), dtype=np.float32)
|
|
_pred[mask] = pred
|
|
c_loss, c_ratio = collision_loss(torch.from_numpy(_pred).unsqueeze(0).to(device), torch.from_numpy(body).unsqueeze(0).to(device), torch.from_numpy(mask).unsqueeze(0).to(device), thr=0.0005)
|
|
return c_loss.cpu().numpy()[0], c_ratio.cpu().numpy()
|
|
|
|
def calculate():
|
|
errors_hood = []
|
|
collisions_hood = []
|
|
areas_hood = []
|
|
collision_ratios_hood = []
|
|
if calculate_ours:
|
|
errors_ours = []
|
|
collisions_ours = []
|
|
areas_ours = []
|
|
collision_ratios_ours = []
|
|
|
|
garments = sorted(os.listdir(hood_folder))[args.start:args.end]
|
|
garments = [g[:-4] for g in garments]
|
|
body_faces = np.load('smpl/smpl_faces.npy').astype(np.int32)
|
|
|
|
for garment in tqdm(garments):
|
|
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:]
|
|
template, template_faces, _, _ = readOBJ(os.path.join(cloth3d_folder, id, garment_type + '.obj'))
|
|
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']
|
|
|
|
m_template = tm.Trimesh(vertices=template, faces=template_faces, process=False, validate=False)
|
|
|
|
for i in range(pred_hood.shape[0]):
|
|
gt_tree = cKDTree(gt[i])
|
|
|
|
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']
|
|
pred_ours = np.load(f'{ours_folder[picker]}/{id}_{i+2}_{garment_type}.npz', allow_pickle=True)['arr_0']
|
|
except:
|
|
print('missed one')
|
|
continue
|
|
|
|
# m_body = tm.Trimesh(vertices=body[i], faces=body_faces, process=False, validate=False)
|
|
|
|
errors_hood.append((f'{id}_{i}_{garment_type}', chamfer(pred_hood[i], gt[i], gt_tree)))
|
|
m_hood = tm.Trimesh(vertices=pred_hood[i], faces=pred_hood_faces, process=False, validate=False)
|
|
# collision_hood = get_collision_depth(m_body, m_hood)
|
|
# collisions_hood.append((f'{id}_{i}_{garment_type}', collision_hood))
|
|
area_hood = get_area_diff(m_template, m_hood)
|
|
areas_hood.append((f'{id}_{i}_{garment_type}', area_hood))
|
|
collision_ratio_hood = get_collision_loss(pred_hood[i], body[i], mask)[1]
|
|
collision_ratios_hood.append((f'{id}_{i}_{garment_type}', collision_ratio_hood))
|
|
|
|
if calculate_ours:
|
|
errors_ours.append((f'{id}_{i}_{garment_type}', chamfer(pred_ours[mask], gt[i], gt_tree)))
|
|
m_ours = tm.Trimesh(vertices=pred_ours[mask], faces=faces, process=False, validate=False)
|
|
# collision_ours = get_collision_depth(m_body, m_ours)
|
|
# collisions_ours.append((f'{id}_{i}_{garment_type}', collision_ours))
|
|
area_ours = get_area_diff(m_template, m_ours)
|
|
areas_ours.append((f'{id}_{i}_{garment_type}', area_ours))
|
|
collision_ratio_ours = get_collision_loss(pred_ours[mask], body[i], mask)[1]
|
|
# collision_ratio_ours_2 = get_collision_loss_2(pred_ours[mask], body[i], mask)[1]
|
|
collision_ratios_ours.append((f'{id}_{i}_{garment_type}', collision_ratio_ours))
|
|
|
|
|
|
# fig = go.Figure(data=[go.Mesh3d(x=pred_ours[mask][:,0], y=pred_ours[mask][:,2], z=pred_ours[mask][:,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=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()
|
|
|
|
# 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()
|
|
# pass
|
|
|
|
# fig = go.Figure()
|
|
# fig.add_trace(go.Scatter3d(x=gt[i][:, 0], y=gt[i][:, 2], z=gt[i][:, 1], mode='markers', marker=dict(size=1)))
|
|
# fig.add_trace(go.Scatter3d(x=pred_hood[i][:, 0], y=pred_hood[i][:, 2], z=pred_hood[i][:, 1], mode='markers', marker=dict(size=1)))
|
|
# fig.update_layout(scene_aspectmode='cube')
|
|
# fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1]), yaxis=dict(range=[-1, 1]), zaxis=dict(range=[-1, 1])))
|
|
# fig.show()
|
|
# pass
|
|
|
|
# fig = go.Figure()
|
|
# fig.add_trace(go.Scatter3d(x=gt[i][:, 0], y=gt[i][:, 2], z=gt[i][:, 1], mode='markers', marker=dict(size=1)))
|
|
# fig.add_trace(go.Scatter3d(x=pred_ours[i,mask,:][:, 0], y=pred_ours[i,mask,:][:, 2], z=pred_ours[i,mask,:][:, 1], mode='markers', marker=dict(size=1)))
|
|
# fig.update_layout(scene_aspectmode='cube')
|
|
# fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1]), yaxis=dict(range=[-1, 1]), zaxis=dict(range=[-1, 1])))
|
|
# fig.show()
|
|
|
|
|
|
# save errors to a csv file
|
|
os.makedirs(f'{results_folder}/hood/{name}_hood', exist_ok=True)
|
|
os.makedirs(f'{results_folder}/hood/{name}_hood/chamfer', exist_ok=True)
|
|
os.makedirs(f'{results_folder}/hood/{name}_hood/collision', exist_ok=True)
|
|
os.makedirs(f'{results_folder}/hood/{name}_hood/area', exist_ok=True)
|
|
os.makedirs(f'{results_folder}/hood/{name}_hood/collision_ratio', exist_ok=True)
|
|
|
|
for frame, error in errors_hood:
|
|
np.save(f'{results_folder}/hood/{name}_hood/chamfer/{frame}.npy', error)
|
|
for frame, error in collisions_hood:
|
|
np.save(f'{results_folder}/hood/{name}_hood/collision/{frame}.npy', error)
|
|
for frame, error in areas_hood:
|
|
np.save(f'{results_folder}/hood/{name}_hood/area/{frame}.npy', error)
|
|
for frame, error in collision_ratios_hood:
|
|
np.save(f'{results_folder}/hood/{name}_hood/collision_ratio/{frame}.npy', error)
|
|
|
|
os.makedirs(f'{results_folder}/hood/{name}_ours', exist_ok=True)
|
|
os.makedirs(f'{results_folder}/hood/{name}_ours/chamfer', exist_ok=True)
|
|
os.makedirs(f'{results_folder}/hood/{name}_ours/collision', exist_ok=True)
|
|
os.makedirs(f'{results_folder}/hood/{name}_ours/area', exist_ok=True)
|
|
os.makedirs(f'{results_folder}/hood/{name}_ours/collision_ratio', exist_ok=True)
|
|
|
|
for frame, error in errors_ours:
|
|
np.save(f'{results_folder}/hood/{name}_ours/chamfer/{frame}.npy', error)
|
|
for frame, error in collisions_ours:
|
|
np.save(f'{results_folder}/hood/{name}_ours/collision/{frame}.npy', error)
|
|
for frame, error in areas_ours:
|
|
np.save(f'{results_folder}/hood/{name}_ours/area/{frame}.npy', error)
|
|
for frame, error in collision_ratios_ours:
|
|
np.save(f'{results_folder}/hood/{name}_ours/collision_ratio/{frame}.npy', error)
|
|
|
|
# with open(f'{results_folder}/hood/{name}/chamfer_hood.csv', 'w') as f:
|
|
# for frame, error in errors_hood:
|
|
# frame_str = ','.join(str(frame).split('_'))
|
|
# f.write(f'{frame_str},{error}\n')
|
|
|
|
# with open(f'{results_folder}/hood/{name}/collision_hood.csv', 'w') as f:
|
|
# for frame, error in collisions_hood:
|
|
# frame_str = ','.join(str(frame).split('_'))
|
|
# f.write(f'{frame_str},{error}\n')
|
|
|
|
# with open(f'{results_folder}/hood/{name}/area_hood.csv', 'w') as f:
|
|
# for frame, error in areas_hood:
|
|
# frame_str = ','.join(str(frame).split('_'))
|
|
# f.write(f'{frame_str},{error}\n')
|
|
|
|
# if calculate_ours:
|
|
# with open(f'{results_folder}/hood/{name}_chamfer_ours.csv', 'w') as f:
|
|
# for frame, error in errors_ours:
|
|
# frame_str = ','.join(str(frame).split('_'))
|
|
# f.write(f'{frame_str},{error}\n')
|
|
|
|
# with open(f'{results_folder}/hood/{name}/collision_ours.csv', 'w') as f:
|
|
# for frame, error in collisions_ours:
|
|
# frame_str = ','.join(str(frame).split('_'))
|
|
# f.write(f'{frame_str},{error}\n')
|
|
|
|
# with open(f'{results_folder}/hood/{name}/area_ours.csv', 'w') as f:
|
|
# for frame, error in areas_ours:
|
|
# frame_str = ','.join(str(frame).split('_'))
|
|
# f.write(f'{frame_str},{error}\n')
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
|
name = 'hood_3_var'
|
|
|
|
hood_folder = '/data-local/127_2/results/hood/all_3'
|
|
ours_folder = dict()
|
|
ours_folder['skirt'] = '/data-local/127_2/results/pyr_iccvw_64_128_256_512_skirt_e2e_var_3_450000'
|
|
ours_folder['nonskirt'] = '/data-local/127_2/results/pyr_iccvw_64_128_256_512_e2e_var_2_600000'
|
|
cloth3d_folder = '/data-net/dataset/cloth3d/test'
|
|
|
|
results_folder = '/data-local/127_2/results'
|
|
|
|
calculate_ours = True
|
|
|
|
|
|
|
|
calculate() |