new scripts
This commit is contained in:
@@ -0,0 +1,261 @@
|
||||
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')
|
||||
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))
|
||||
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:
|
||||
for i in range(4):
|
||||
_pred = pred.item()[f'lvl{i}']
|
||||
c_loss, c_ratio = get_collision_loss(_pred, body, 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()
|
||||
|
||||
|
||||
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)))
|
||||
|
||||
|
||||
# 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()
|
||||
Reference in New Issue
Block a user