152 lines
5.3 KiB
Python
152 lines
5.3 KiB
Python
import os
|
|
import sys
|
|
from configs.config_abl_all_post 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
|
|
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 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())
|
|
|
|
# smpl = SMPL(use_one_hot=config['use_one_hot'])
|
|
skirt_modifier = '_skirt' if config['is_skirt'] else ''
|
|
|
|
folder = config['test_dir']
|
|
|
|
checkpoint_name = config['name']
|
|
iteration = ''
|
|
|
|
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/data/results/post/' + checkpoint_name + '-' + iteration
|
|
|
|
|
|
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 }
|
|
|
|
masks_folder = '/home/workspace/data/dataset/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')]
|
|
|
|
if config['remove_global_orientation']:
|
|
gt_name = 'mesh_posed_no_orient'
|
|
else:
|
|
gt_name = 'mesh_posed'
|
|
|
|
orig_folder = '/home/workspace/data/dataset/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.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)))
|
|
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('/home/workspace/data/results/errors/post/' + checkpoint_name + '-' + iteration + '.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()
|
|
|