194 lines
7.7 KiB
Python
194 lines
7.7 KiB
Python
import os
|
|
import sys
|
|
from configs.iccv.config_abl_pyr_32_64_128_tshirt import config
|
|
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
|
|
import tensorflow as tf
|
|
|
|
gpus = tf.config.experimental.list_physical_devices('GPU')
|
|
if gpus:
|
|
try:
|
|
# Currently, memory growth needs to be the same across GPUs
|
|
for gpu in gpus:
|
|
tf.config.experimental.set_memory_growth(gpu, True)
|
|
logical_gpus = tf.config.experimental.list_logical_devices('GPU')
|
|
print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
|
|
except RuntimeError as e:
|
|
# Memory growth must be set before GPUs have been initialized
|
|
print(e)
|
|
|
|
import csv
|
|
import glob
|
|
import time
|
|
import argparse
|
|
import wandb
|
|
from tqdm import tqdm
|
|
import random
|
|
import numpy as np
|
|
import plotly.graph_objects as go
|
|
import plotly.express as ex
|
|
from losses import collision_loss, normal_loss
|
|
if config['is_skirt']:
|
|
from losses import edge_loss_skirt as edge_loss
|
|
else:
|
|
from losses import edge_loss
|
|
|
|
|
|
|
|
# from postprocessing import fix_collision
|
|
# from losses import collision_ratio, kl_loss #, concept_loss
|
|
# # from losses import collision_penalty as collision_penalty
|
|
# from losses import collision_loss as collision_penalty
|
|
# from custom_utils.custom_utils import batch_to_dict, load_model, plot_3d, UV2Mesh, get_summary_html, maskFacesBatch, _maskFaces, plot_3d_collision
|
|
# from custom_utils.smpl import SMPL
|
|
# from custom_utils.supersmpl_io import writeOBJ
|
|
|
|
|
|
|
|
|
|
# uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_adaptive=False)
|
|
|
|
if __name__ == '__main__':
|
|
parser = argparse.ArgumentParser(description='')
|
|
parser.add_argument('--name', type=str)
|
|
args = vars(parser.parse_args())
|
|
|
|
print(config['name'])
|
|
|
|
# smpl = SMPL(use_one_hot=config['use_one_hot'])
|
|
skirt_modifier = '_skirt' if config['is_skirt'] else ''
|
|
|
|
error_type = 'euclidean' # 'euclidean' 'collision' edge normal
|
|
|
|
folder = config['test_dir']
|
|
|
|
checkpoint_name = config['name']
|
|
# iteration = '200000'
|
|
|
|
F = np.load("custom_utils/smpl_faces" + skirt_modifier + ".npy").astype(np.int32)
|
|
garment_type_table = np.load('custom_data/garment_type_table_test.npy', allow_pickle=True).item()
|
|
|
|
|
|
output_folder = '/home/workspace/local-data/results/iccv/' + checkpoint_name + '-'
|
|
# check if there is a folder that starts with this name using glob
|
|
folders = glob.glob(output_folder + '*')
|
|
assert len(folders) == 1, f'There should be only one checkpoint that starts with this name, but found {len(folders)}'
|
|
output_folder = folders[0]
|
|
iteration = output_folder.split('-')[-1]
|
|
|
|
|
|
# faces_folder = '/home/workspace/data/results/support/faces' + ('_skirt' if config['is_skirt'] else '')
|
|
# ids = os.listdir(faces_folder)
|
|
# ids = [id.split('.')[0] for id in ids]
|
|
# faces = {id : np.load(faces_folder + '/' + id + '.npy') for id in ids }
|
|
|
|
print('Reading masks...')
|
|
masks_folder = '/home/workspace/local-data' + ('' if config['is_skirt'] else '_127') + '/test' + ('_skirt' if config['is_skirt'] else '') + '_ready' + ('_2' if config['is_skirt'] else '') + '/mesh_mask'
|
|
ids = os.listdir(masks_folder)
|
|
ids = [id.split('.')[0] for id in ids]
|
|
masks = {id : np.load(masks_folder + '/' + id + '.npy') for id in ids }
|
|
|
|
frames = os.listdir(output_folder)
|
|
frames = [frame for frame in frames if not frame.endswith('_fixed_coll.npy')]
|
|
# frames = frames[:100]
|
|
|
|
# if 'frames_file_test' in config.keys():
|
|
# with open(os.path.join(folder, config['frames_file_test']), 'r') as f:
|
|
# lines = f.readlines()
|
|
# lines = [line.strip() for line in lines]
|
|
# frames = lines
|
|
|
|
if error_type == 'collision':
|
|
print('Reading bodies...')
|
|
bodies_folder = '/home/workspace/local-data' + ('' if config['is_skirt'] else '_127') + '/test' + ('_skirt' if config['is_skirt'] else '') + '_ready' + ('_2' if config['is_skirt'] else '') + '/mesh_posed_body_no_orient'
|
|
bodies = {frame : np.load(bodies_folder + '/' + frame[:-1] + 'z')['mesh_posed_body_no_orient'] for frame in tqdm(frames) }
|
|
|
|
if config['remove_global_orientation']:
|
|
gt_name = 'mesh_posed_no_orient'
|
|
else:
|
|
gt_name = 'mesh_posed'
|
|
|
|
orig_folder = '/home/workspace/local-data' + ('' if config['is_skirt'] else '_127') + '/test' + ('_skirt' if config['is_skirt'] else '') + '_ready' + ('_2' if config['is_skirt'] else '') + '/' + gt_name
|
|
|
|
errors = dict()
|
|
if config['is_skirt']:
|
|
per_garment_errors = { 'Tshirt': [], 'Dress': [], 'Skirt': [], 'Top': [] }
|
|
else:
|
|
per_garment_errors = { 'Tshirt': [], 'Jumpsuit': [], 'Trousers': [], 'Top': [] }
|
|
|
|
for frame in tqdm(frames):
|
|
id = frame.split('_')[0]
|
|
frame_nr = frame.split('.')[0].split('_')[1]
|
|
gt = np.load(os.path.join(orig_folder, id + '_' + frame_nr + '.npz'))[gt_name][masks[id], :].astype(np.float32)
|
|
pred = np.load(os.path.join(output_folder, id + '_' + frame_nr + '.npy'))
|
|
if error_type == 'edge':
|
|
pred_unmasked = np.zeros((1, 14475, 3), dtype=np.float32)
|
|
pred_unmasked[0, masks[id], :] = pred
|
|
gt_unmasked = np.zeros((1, 14475, 3), dtype=np.float32)
|
|
gt_unmasked[0, masks[id], :] = gt
|
|
metric = edge_loss(gt_unmasked, pred_unmasked, masks[id][None,:].astype(np.float32)).numpy()[0]
|
|
elif error_type == 'euclidean':
|
|
metric = np.mean(np.sqrt(np.sum((gt - pred) ** 2, axis=-1)))
|
|
elif error_type == 'collision':
|
|
pred_unmasked = np.zeros((1, 14475, 3), dtype=np.float32)
|
|
pred_unmasked[0, masks[id], :] = pred
|
|
metric = collision_loss(pred_unmasked, bodies[frame][None,:,:], masks[id][None,:].astype(np.float32))[0].numpy()[0]
|
|
elif error_type == 'normal':
|
|
pred_unmasked = np.zeros((1, 14475, 3), dtype=np.float32)
|
|
pred_unmasked[0, masks[id], :] = pred
|
|
gt_unmasked = np.zeros((1, 14475, 3), dtype=np.float32)
|
|
gt_unmasked[0, masks[id], :] = gt
|
|
metric = normal_loss(gt_unmasked, pred_unmasked, masks[id][None,:].astype(np.float32)).numpy()
|
|
else:
|
|
raise ValueError('Error type not supported')
|
|
|
|
errors[id + '_' + frame_nr] = metric
|
|
per_garment_errors[garment_type_table[id]].append(metric)
|
|
|
|
|
|
for key in per_garment_errors.keys():
|
|
print(key + ': ' + str(np.mean(per_garment_errors[key])))
|
|
|
|
with open(f'/home/workspace/data/results/errors/iccv/{checkpoint_name}-{iteration}-{error_type}.csv', mode='w') as f:
|
|
file_writer = csv.writer(f, delimiter=',', quotechar='"', quoting=csv.QUOTE_MINIMAL)
|
|
|
|
for key in sorted(errors.keys()):
|
|
file_writer.writerow([key, errors[key]])
|
|
|
|
# x = gt
|
|
# y = pred
|
|
# fig = go.Figure(data=[
|
|
# go.Mesh3d(
|
|
# x=x[:,0],
|
|
# y=x[:,1],
|
|
# z=x[:,2],
|
|
# i = faces[id][:,0],
|
|
# j = faces[id][:,1],
|
|
# k = faces[id][:,2],
|
|
# ),
|
|
# go.Mesh3d(
|
|
# x=y[:,0],
|
|
# y=y[:,1],
|
|
# z=y[:,2],
|
|
# i = faces[id][:,0],
|
|
# j = faces[id][:,1],
|
|
# k = faces[id][:,2],
|
|
# ),
|
|
# ])
|
|
|
|
# fig.update_layout(scene_aspectmode='manual',
|
|
# scene_aspectratio=dict(x=1, y=1, z=1))
|
|
# fig.update_layout(
|
|
# scene = dict(
|
|
# xaxis = dict(nticks=4, range=[-1,1],),
|
|
# yaxis = dict(nticks=4, range=[-1,1],),
|
|
# zaxis = dict(nticks=4, range=[-1,1],),)
|
|
# )
|
|
|
|
# fig.show()
|
|
|
|
# print()
|
|
|
|
# print()
|
|
|