176 lines
6.6 KiB
Python
176 lines
6.6 KiB
Python
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='')
|
|
parser.add_argument('-g', '--gpu', type=str)
|
|
args = vars(parser.parse_args())
|
|
|
|
if args['debug']:
|
|
from configs.iccvw.pyr_64_128_cape_f_2 import 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"] = config_1['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 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
|
|
|
|
|
|
|
|
def calculate():
|
|
errors = []
|
|
|
|
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')['arr_0']
|
|
gt = readPC2Frame(f'{config_1["test_dir"]}/mesh_posed_no_orient/{frame.id}_{frame.garment}.pc16', frame.frame_nr, True)
|
|
# mask = np.load(f'{config["test_dir"]}/mesh_mask/{frame.id}_{frame.garment}.npy')
|
|
|
|
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']
|
|
|
|
pred = pred[new_mask]
|
|
gt = gt[mask_convert]
|
|
|
|
error = np.linalg.norm(pred - gt, axis=1).mean()
|
|
errors.append((str(frame), error))
|
|
|
|
# save errors to a csv file
|
|
with open(f'{results_folder}/errors/{name}.csv', 'w') as f:
|
|
for frame, error in errors:
|
|
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() |