import os import sys from configs.iccv.config_abl_pyr_64_cape_m_00032_SL 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 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 from PIL import Image from postprocessing import fix_collision if config['is_skirt']: from losses import edge_loss_skirt as edge_loss else: from losses import edge_loss 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 dataset import DatasetFactory, StaticDataset from custom_utils.smpl import SMPL from models.pyramid_model_optimized import PyramidNew as Pyramid from custom_utils.supersmpl_io import writeOBJ from unif_metric import chamfer_and_score uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_adaptive=False, is_smpl_orig=config['is_cape']) def group_inputs_for_levels(batch): levels_input = [] for res in config['pyramid']['resolutions']: res_modifier = '_' + str(res) if res != 512 else '' levels_input.append([batch['uv_unposed' + res_modifier], batch['uv_static' + res_modifier], batch['uv_body_posed' + res_modifier], batch['uv_curvature' + res_modifier], batch['uv_mask' + res_modifier], batch['uv_mask']]) conds = [batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature']] return levels_input, conds @tf.function def step(batch): levels_input, conds = group_inputs_for_levels(batch) out_levels = model([levels_input, conds], training=False) pred = uv2mesh(out_levels[-1][0]) mesh_pred = uv2mesh(pred) # if config['use_skinning']: # shape = batch['pose_shape'][:, 72:] # gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32) # shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])]) # mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis] # mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped) # else: # mesh_pred_posed = mesh_pred # mesh_pred_posed_fixed = fix_collision(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']) # return mesh_pred_posed, None #mesh_pred_posed_fixed return mesh_pred if __name__ == '__main__': # init_wandb() parser = argparse.ArgumentParser(description='') parser.add_argument('--name', type=str) args = vars(parser.parse_args()) skirt_modifier = '_skirt' if config['is_skirt'] else '' folder = config['test_dir'] smpl = SMPL(use_one_hot=False, orig_smpl=config['is_cape']) model = Pyramid(config) # str_split = config['name'].split('_') # test_idx = str_split.index('test') # checkpoint_name = '_'.join(str_split[:test_idx]) checkpoint_name = config['name'] if config['is_cape']: F = np.load("custom_utils/smpl_faces_orig.npy").astype(np.int32) else: F = np.load("custom_utils/smpl_faces" + skirt_modifier + ".npy").astype(np.int32) # ids = os.listdir(os.path.join(folder, 'pose_shape')) # ids = [id.split('_')[0] for id in ids] # ids = set(ids) iteration = '' if iteration == '': iteration = load_model(model, 'checkpoints/iccv/' + checkpoint_name, return_iteration=True) else: load_model(model, 'checkpoints/iccv/' + checkpoint_name, iteration=iteration) output_folder = '/home/workspace/data/results/' + checkpoint_name + '-' + iteration # os.makedirs(output_folder, exist_ok=True) # if config['use_skinning']: # config['signature_dynamic'].pop('mesh_posed') # config['signature_dynamic'].pop('mesh_unposed') # config['signature_dynamic'].pop('mesh_body_posed') dataset_resolutions = config['pyramid']['resolutions'].copy() if 512 not in dataset_resolutions: dataset_resolutions.append(512) factory = DatasetFactory() # dataset_test = factory.get_dataset(folder=config['test_dir'], # batch_size=config['batch_size'], # signature=config['signature_dynamic'], # resolutions=[config['resolution']], # in_memory=False, # do_repeat=False, # use_no_orient=config['remove_global_orientation'], # frames_file=config['frames_file_test'] if 'frames_file_test' in config.keys() else None, # is_cape=config['is_cape']) # static_dataset_test = StaticDataset(folder=config['test_dir'], # signature=config['signature_static'], # resolutions=[config['resolution']], # in_memory=True, # is_cape=config['is_cape']) dataset_test = factory.get_dataset(folder=config['test_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], resolutions=dataset_resolutions, in_memory=False, do_repeat=False, use_no_orient=True, frames_file=config['frames_file_test'] if 'frames_file_test' in config.keys() else None, is_cape=config['is_cape']) static_dataset_test = StaticDataset(folder=config['test_dir'], signature=config['signature_static'], resolutions=dataset_resolutions, in_memory=True, is_cape=config['is_cape']) # 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 } errors = [] euclid = [] for batch in tqdm.tqdm(dataset_test): batch = static_dataset_test.expand_batch(batch_to_dict(batch, config['signature_dynamic'])) mesh_pred = step(batch) # mesh_pred_posed, mesh_pred_posed_fixed = mesh_pred_posed.numpy(), mesh_pred_posed_fixed.numpy() mask = batch['mesh_mask'].numpy().astype(bool) for i in range(mask.shape[0]): id, frame_nr = batch['id'][i].numpy().decode("utf-8"), batch['frame_nr'][i].numpy().decode("utf-8") x = chamfer_and_score(mesh_pred.numpy()[i, mask[i], :], batch['mesh_unposed'].numpy()[i, mask[i], :]) y = np.mean(np.sqrt(np.sum((mesh_pred.numpy()[i, mask[i], :] - batch['mesh_unposed'].numpy()[i, mask[i], :])**2, axis=-1))) errors.append(x[1]) euclid.append(y) # writeOBJ(os.path.join(output_folder, id + '_' + frame_nr + '.obj'), mesh_pred_posed[i, mask[i], :], faces[id]) # np.save(os.path.join(output_folder, id + '_' + frame_nr + '.npy'), mesh_pred_posed[i, mask[i], :]) # np.save(os.path.join(output_folder, id + '_' + frame_nr + '_fixed_coll.npy'), mesh_pred_posed_fixed[i, mask[i], :]) # writeOBJ(os.path.join(output_folder, id + '_' + frame_nr + '_fixed_coll.obj'), mesh_pred_posed_fixed[i, mask[i], :], faces[id]) print(np.mean(errors)) print(np.mean(euclid)) print()