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()