import os import sys from configs.iccv.config_abl_pyr_32_64_128_skin_skirt 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 = 'collision' # '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] 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()