import os import sys from configs.config_abl_overfit_jumpsuit_2 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.cvae_model import CVAE 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']) @tf.function def step(batch): if config['use_skinning']: uv_orig = batch['uv_unposed'] else: uv_orig = batch['uv_posed'] pred, _, _ = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False) 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 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 = CVAE(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 = '' load_model(model, 'checkpoints/' + checkpoint_name, iteration=iteration) # load_model(model, 'checkpoints/' + 'abl_all_120k') output_folder = '/home/workspace/data/results/post/' + 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') 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']) # 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 } for batch in tqdm.tqdm(dataset_test): batch = static_dataset_test.expand_batch(batch_to_dict(batch, config['signature_dynamic'])) mesh_pred_posed, mesh_pred_posed_fixed = step(batch) # mesh_pred_posed, mesh_pred_posed_fixed = mesh_pred_posed.numpy(), mesh_pred_posed_fixed.numpy() 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])