This commit is contained in:
2023-03-10 17:25:44 +00:00
parent cb76d36c4f
commit 4b1cb0b06d
27 changed files with 1890 additions and 207 deletions
+53 -23
View File
@@ -1,37 +1,37 @@
import os
import sys
from configs.iccv.config_abl_pyr_32_64_25p import config
from configs.iccv.config_abl_pyr_32_64_128_skin_skirt import config
os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
# import tensorflow as tf
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)
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
import tqdm
from tqdm import tqdm
import random
import numpy as np
import plotly.graph_objects as go
import plotly.express as ex
# if config['is_skirt']:
# from losses import edge_loss_skirt as edge_loss
# else:
# from losses import edge_loss
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
@@ -53,9 +53,13 @@ if __name__ == '__main__':
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 = 'collision' # 'euclidean' 'collision' edge normal
folder = config['test_dir']
checkpoint_name = config['name']
@@ -65,7 +69,7 @@ if __name__ == '__main__':
garment_type_table = np.load('custom_data/garment_type_table_test.npy', allow_pickle=True).item()
output_folder = '/home/workspace/local-data_128/results/iccv/' + checkpoint_name + '-'
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)}'
@@ -78,6 +82,7 @@ if __name__ == '__main__':
# 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]
@@ -85,6 +90,11 @@ if __name__ == '__main__':
frames = os.listdir(output_folder)
frames = [frame for frame in frames if not frame.endswith('_fixed_coll.npy')]
# frames = frames[:100]
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'
@@ -99,12 +109,32 @@ if __name__ == '__main__':
else:
per_garment_errors = { 'Tshirt': [], 'Jumpsuit': [], 'Trousers': [], 'Top': [] }
for frame in tqdm.tqdm(frames):
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'))
metric = np.mean(np.sqrt(np.sum((gt - pred) ** 2, axis=-1)))
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)
@@ -112,7 +142,7 @@ if __name__ == '__main__':
for key in per_garment_errors.keys():
print(key + ': ' + str(np.mean(per_garment_errors[key])))
with open('/home/workspace/data/results/errors/iccv/' + checkpoint_name + '-' + iteration + '.csv', mode='w') as f:
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()):