import os import sys from configs.iccv.config_abl_pyr_64_128_cape_f 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 models.pyramid_model_optimized import PyramidNew as Pyramid 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, group_inputs_for_levels from dataset import DatasetFactory, StaticDataset from custom_utils.smpl import SMPL from models.cvae_model import CVAE from custom_utils.supersmpl_io import writeOBJ uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_adaptive=False) @tf.function def step(batch): levels_input, conds = group_inputs_for_levels(batch, config) out_levels = model([levels_input, conds], training=True) pred_final = out_levels[-1][0] mesh_pred = uv2mesh(pred_final) 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 = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped) # mesh_pred_posed_fixed = fix_collision(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']) return mesh_pred_posed, None 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) model: Pyramid = Pyramid(config) # str_split = config['name'].split('_') # test_idx = str_split.index('test') # checkpoint_name = '_'.join(str_split[:test_idx]) checkpoint_name = config['name'] 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) # load_model(model, 'checkpoints/' + 'abl_all_120k') output_folder = '/home/workspace/local-data_128/results/iccv/' + checkpoint_name + '-' + iteration os.makedirs(output_folder, exist_ok=True) dataset_resolutions = config['pyramid']['resolutions'].copy() if 512 not in dataset_resolutions: dataset_resolutions.append(512) if config['use_skinning']: config['signature_dynamic'].pop('mesh_posed') config['signature_dynamic'].pop('mesh_unposed') config['signature_dynamic'].pop('mesh_body_posed') factory = DatasetFactory() 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) static_dataset_test = StaticDataset(folder=config['test_dir'], signature=config['signature_static'], resolutions=dataset_resolutions, in_memory=True) # faces_folder = '/home/workspace/local-data_128/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 } for batch in tqdm.tqdm(dataset_test): batch = static_dataset_test.expand_batch(batch_to_dict(batch, config['signature_dynamic'])) mesh_pred_posed, _ = step(batch) mesh_pred_posed = mesh_pred_posed.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") # 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]) # x = batch['mesh_posed'].numpy()[i, mask[i], :] # y = mesh_pred_posed[i, mask[i], :] # 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()