import os import sys sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..') import importlib import argparse parser = argparse.ArgumentParser(description='') parser.add_argument('--name', type=str) parser.add_argument('--config', type=str) parser.add_argument('--debug', action='store_true') parser.add_argument('--iteration', type=str, default='200000') parser.add_argument('-g', '--gpu', type=str, default='0') parser.add_argument('-f', '--final', action='store_true') args = vars(parser.parse_args()) if args['debug']: from configs.iccvw.pyr_64_128_256_512_skirt_e2e import config as config_1 else: config_module = importlib.import_module(args['config']) config_1 = config_module.config os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID" os.environ["CUDA_VISIBLE_DEVICES"] = args['gpu'] if args['debug']: os.environ["CUDA_LAUNCH_BLOCKING"] = '1' import time import argparse import wandb import tqdm import random import numpy as np import plotly.graph_objects as go import plotly.express as ex from PIL import Image import trimesh as tm import torch import torch.nn as nn import torch.optim as optim import torch.nn.functional as F from tqdm import tqdm from data.dataset import Dataset from data.utils import Frame from data.cloth3d_io import readPC2Frame, readPC2 from utils.checkpoints import CheckpointSaver, load_model from torch.utils.data import DataLoader from utils.utils import text_to_html from utils.postprocess import fix_collision from losses.losses import collision_loss, normal_loss, kl_loss from custom_utils.custom_utils import UV2Mesh, Accumulator, plot_3d, get_summary_html, maskFacesBatch, get_summary_str, plot_3d_collision, group_inputs_for_levels from plotly.subplots import make_subplots if config_1['is_skirt']: from losses.losses import edge_loss_skirt as edge_loss else: from losses.losses import edge_loss from losses.losses import collision_loss, normal_loss, kl_loss device = 'cuda' def get_edge_loss(pred, gt, mask): e_loss = torch.mean(edge_loss(torch.from_numpy(pred).unsqueeze(0).to(device), torch.from_numpy(gt).unsqueeze(0).to(device), torch.from_numpy(mask).unsqueeze(0).to(device))) return e_loss.cpu().numpy() def get_normal_loss(pred, gt, mask): n_loss = normal_loss(torch.from_numpy(pred).unsqueeze(0).to(device), torch.from_numpy(gt).unsqueeze(0).to(device), torch.from_numpy(mask).unsqueeze(0).to(device), is_skirt=config_1['is_skirt'], is_smpl_orig=config_1['is_cape']) return n_loss.cpu().numpy() def get_collision_loss(pred, body, mask): 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_collision = [] errors_collision_ratio = [] errors_normal = [] errors_edge = [] for frame in tqdm(frames): # import plotly.graph_objects as go # mesh_pred = mesh_pred.cpu().numpy() # mask = batch['mesh_mask'].cpu().numpy() # i = 3 # fig = go.Figure(data=[go.Scatter3d(x=mesh_pred[i,mask[i],0], y=mesh_pred[i,mask[i],2], z=mesh_pred[i,mask[i],1], mode='markers', marker=dict(size=1))]) # fig.update_layout(scene_aspectmode='cube') # fig.show() pred = np.load(f'{out_folder}/{frame}.npz', allow_pickle=True)['arr_0'] try: pred.item()['lvl3'] all_levels = True except: all_levels = False gt = readPC2Frame(f'{config_1["test_dir"]}/mesh_posed_no_orient/{frame.id}_{frame.garment}.pc16', frame.frame_nr, True) body = readPC2Frame(f'{config_1["test_dir"]}/mesh_body_posed_no_orient/{frame.id}.pc16', frame.frame_nr, True) fix_data = np.load(f'{config_1["test_dir"]}/mesh_fixed/{frame.id}_{frame.garment}.npz') mask_convert = fix_data['mas_convert'] new_mask = fix_data['mask'] gt = gt[mask_convert] gt_full = np.zeros((14475, 3), dtype=np.float32) gt_full[new_mask] = gt gt = gt_full if not all_levels: c_loss, c_ratio = get_collision_loss(pred, body, new_mask) errors_collision.append((str(frame), c_loss)) errors_collision_ratio.append((str(frame), c_ratio)) errors_normal.append((str(frame), get_normal_loss(pred, gt, new_mask))) errors_edge.append((str(frame), get_edge_loss(pred, gt, new_mask))) else: levels = [0,1,2,3] if not args['final'] else [3] for i in levels: _pred = pred.item()[f'lvl{i}'] c_loss, c_ratio = get_collision_loss(_pred, body, new_mask) errors_collision.append((f'{frame}_lvl{i}', c_loss)) errors_collision_ratio.append((f'{frame}_lvl{i}', c_ratio)) errors_normal.append((f'{frame}_lvl{i}', get_normal_loss(_pred, gt, new_mask))) errors_edge.append((f'{frame}_lvl{i}', get_edge_loss(_pred, gt, new_mask))) # faces = fix_data['faces'] # body_faces = np.load('smpl/smpl_faces.npy').astype(np.int32) # m1 = tm.Trimesh(vertices=body, faces=body_faces, process=False, validate=False) # m2 = tm.Trimesh(vertices=_pred[new_mask], faces=faces, process=False, validate=False) # # m2 = tm.Trimesh(vertices=gt[new_mask], faces=faces, process=False, validate=False) # 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]) # _pred = gt # fig = go.Figure(data=[go.Mesh3d(x=_pred[new_mask,0], y=_pred[new_mask,2], z=_pred[new_mask,1], i=faces[:,0], j=faces[:,1], k=faces[:,2])]) # fig.add_trace(go.Mesh3d(x=body[:,0], y=body[:,2], z=body[:,1], i=body_faces[:,0], j=body_faces[:,1], k=body_faces[:,2])) # 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() # save errors to a csv file with open(f'{results_folder}/errors/{name}_collision.csv', 'w') as f: for frame, error in errors_collision: frame_str = ','.join(str(frame).split('_')) f.write(f'{frame_str},{error}\n') with open(f'{results_folder}/errors/{name}_collision_ratio.csv', 'w') as f: for frame, error in errors_collision_ratio: frame_str = ','.join(str(frame).split('_')) f.write(f'{frame_str},{error}\n') with open(f'{results_folder}/errors/{name}_normal.csv', 'w') as f: for frame, error in errors_normal: frame_str = ','.join(str(frame).split('_')) f.write(f'{frame_str},{error}\n') with open(f'{results_folder}/errors/{name}_edge.csv', 'w') as f: for frame, error in errors_edge: frame_str = ','.join(str(frame).split('_')) f.write(f'{frame_str},{error}\n') def calculate_cape(): from smpl.smpl_torch import SMPL errors = [] clothing_verts_idx = np.load(os.path.join('custom_data', 'clothing_verts_idx.npy')) cape_mask = np.zeros((6890), bool) cape_mask[clothing_verts_idx] = True smpl = SMPL(path='smpl', orig_smpl=True, device='cpu') for frame in tqdm(frames): # import plotly.graph_objects as go # mesh_pred = mesh_pred.cpu().numpy() # mask = batch['mesh_mask'].cpu().numpy() # i = 3 # fig = go.Figure(data=[go.Scatter3d(x=mesh_pred[i,mask[i],0], y=mesh_pred[i,mask[i],2], z=mesh_pred[i,mask[i],1], mode='markers', marker=dict(size=1))]) # fig.update_layout(scene_aspectmode='cube') # fig.show() if not os.path.exists(f'{out_folder}/{frame}.npz'): print('missed one') continue pred = np.load(f'{out_folder}/{frame}.npz')['arr_0'] gt = np.load(f'{config_1["test_dir"]}/mesh_unposed_no_shape/{frame.id}.{frame.frame_nr}.npz')['mesh_unposed_no_shape'] pose_shape = np.load(f'{config_1["test_dir"]}/pose_shape/{frame.id}.{frame.frame_nr}.npy').astype(np.float32) pose, shape = pose_shape[:72], pose_shape[72:] gender = int(shape[10] >= 0.5) shape_offsets = torch.tensordot(torch.from_numpy(shape[:10][None, :]), smpl.shapedirs[gender], dims=[[1], [2]])[0,smpl.to_keep,:].numpy() gt_shaped = gt + shape_offsets min_body = np.load(os.path.join('/data-local/127_2/cape_processed', 'minimal_body_shape', frame.id.split('_')[0], frame.id.split('_')[0] + '_minimal.npy')) offsets_pred = pred[cape_mask[smpl.to_keep],:] - min_body[clothing_verts_idx,:] offsets_orig = gt_shaped[cape_mask[smpl.to_keep],:] - min_body[clothing_verts_idx,:] error = np.linalg.norm(offsets_pred - offsets_orig, axis=1).mean() errors.append((str(frame), error)) # import plotly.graph_objects as go # fig = go.Figure(data=[go.Scatter3d(x=pred[:,0], y=pred[:,2], z=pred[:,1], mode='markers', marker=dict(size=1))]) # fig.add_trace(go.Scatter3d(x=gt[:,0], y=gt[:,2], z=gt[:,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 with open(f'{results_folder}/errors/{name}.csv', 'w') as f: for frame, error in errors: split = str(frame).split('_') if not config_1['is_cape']: frame_str = ','.join(split) else: id = '_'.join(split[:-2]) frame_str = f'{id},{split[-2]},{split[-1]}' f.write(f'{frame_str},{error}\n') if __name__ == '__main__': skirt_modifier = '_skirt' if config_1['is_skirt'] else '' name = config_1['name'] results_folder = '/data-local/127_2/results' if args['iteration'] == '': folders = [f for f in os.listdir(results_folder) if os.path.isdir(os.path.join(results_folder, f)) and f.startswith(config_1['name'])] assert len(folders) == 1 name = folders[0] else: name = f'{name}_{args["iteration"]}' out_folder = f'{results_folder}/{name}' os.makedirs(f'{results_folder}/errors', exist_ok=True) if config_1['is_cape']: from data.dataset_cape import read_frames frames = read_frames(config_1['test_dir'], shuffle=False, is_cape=True) calculate_cape() else: frames_file = config_1['test_frames'] with open(os.path.join(config_1['test_dir'], frames_file), 'r') as f: frames = f.readlines() frames = [Frame(frames.split()[0], int(frames.split()[1]), frames.split()[2]) for frames in frames] calculate()