luni/cvpr
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

save

Browse Source
This commit is contained in:
luni
2023-03-12 23:29:47 +00:00
parent 4b1cb0b06d
commit d5bade988f
17 changed files with 2550 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cvae_pyramid.py
+30 -11
View File
@@ -1,6 +1,6 @@
import os
import sys
from configs.iccv.config_abl_pyr_32_tshirt import config
from configs.iccv.config_abl_pyr_32_64_128_skirt_dress import config
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
@@ -39,12 +39,13 @@ else:
from losses import uv_mask_loss, bend_loss, kl_loss, collision_ratio, normal_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
from custom_utils.custom_utils import batch_to_dict, load_model, save_stl
from dataset import DatasetFactory, StaticDataset
from custom_utils.smpl import SMPL
from models.pyramid_model_optimized import PyramidNew as Pyramid
from postprocessing import fix_collision
from custom_utils.custom_utils import Accumulator
from sample_pro_2 import postprocess
def init_wandb():
@@ -682,9 +683,10 @@ def test():
def generate_results():
load_model(model, 'checkpoints/' + config['name'])
load_model(model, 'checkpoints/iccv/' + config['name'])
F = np.load("custom_utils/smpl_faces" + skirt_modifier + ".npy").astype(np.int32)
dataset_test = factory.get_dataset(folder=config['test_dir'],
dataset_test = factory.get_dataset(folder=config['train_dir'],
batch_size=config['batch_size'],
# batch_size=1,
signature=config['signature_dynamic'],
@@ -693,15 +695,16 @@ def generate_results():
is_cape=config['is_cape'],
frames_file=config['frames_file'] if 'frames_file' in config.keys() else None,
do_repeat=False,
shuffle=False)
static_dataset_test = StaticDataset(folder=config['test_dir'],
shuffle=True)
static_dataset_test = StaticDataset(folder=config['train_dir'],
signature=config['signature_static'],
resolutions=dataset_resolutions,
in_memory=True,
is_cape=config['is_cape'])
folder = os.path.join('/home/workspace/data/results/cape', config['name'])
folder = os.path.join('examples/qualitative', config['name'])
os.makedirs(folder, exist_ok=True)
for batch in tqdm.tqdm(dataset_test):
batch = static_dataset_test.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
@@ -718,14 +721,28 @@ def generate_results():
mesh_pred = uv2mesh(pred)
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)
mesh_gt = batch['mesh_posed'].numpy()
body = batch['mesh_body_posed'].numpy()
for i in range(batch['id'].shape[0]):
mask = batch['mesh_mask'][i].numpy().astype(bool)
id, frame_nr = batch['id'][i].numpy().decode("utf-8"), batch['frame_nr'][i].numpy().decode("utf-8")
unposed = mesh_pred[i].numpy()[mask, :]
posed = mesh_pred_posed[i].numpy()[mask, :]
np.save(os.path.join(folder, id + '_' + frame_nr + '_posed.npy'), posed)
np.save(os.path.join(folder, id + '_' + frame_nr + '_unposed.npy'), unposed)
_, faces_garment = _maskFaces(np.zeros((14475, 3), np.float32), F, mask)
faces_garment = faces_garment.astype(np.int32)
posed_unmasked = np.zeros((1, 14475, 3), np.float32)
posed_unmasked[0, mask, :] = posed
posed = postprocess(posed_unmasked, body[i], faces_garment, mask, do_smoothing=True)
unposed_unmasked = np.zeros((1, 14475, 3), np.float32)
unposed_unmasked[0, mask, :] = unposed
unposed = postprocess(unposed_unmasked, body[i], faces_garment, mask, do_smoothing=True)
save_stl(os.path.join(folder, id + '_' + frame_nr + '_posed.stl'), posed, faces_garment)
save_stl(os.path.join(folder, id + '_' + frame_nr + '_unposed.stl'), unposed, faces_garment)
save_stl(os.path.join(folder, id + '_' + frame_nr + '_gt.stl'), mesh_gt[i][mask, :], faces_garment)
print(id + '_' + frame_nr)
# np.save(os.path.join(folder, id + '_' + frame_nr + '_posed.npy'), posed)
# np.save(os.path.join(folder, id + '_' + frame_nr + '_unposed.npy'), unposed)
pass
@@ -772,7 +789,7 @@ if __name__ == '__main__':
in_memory=config['dataset_in_memory'],
is_cape=config['is_cape'],
frames_file=config['frames_file'] if 'frames_file' in config.keys() else None,
frames_file_skirt=config['frames_file'] if 'frames_file' in config.keys() else '',
frames_file_skirt=config['frames_file_skirt'] if 'frames_file_skirt' in config.keys() else '',
folder_skirt=config['train_dir_skirt'] if 'train_dir_skirt' in config.keys() else '',
)
static_dataset_train = StaticDataset(folder=config['train_dir'],
@@ -788,7 +805,7 @@ if __name__ == '__main__':
in_memory=config['dataset_in_memory'],
is_cape=config['is_cape'],
frames_file=config['frames_file_valid'] if 'frames_file_valid' in config.keys() else None,
frames_file_skirt=config['frames_file_valid'] if 'frames_file_valid' in config.keys() else '',
frames_file_skirt=config['frames_file_valid_skirt'] if 'frames_file_valid_skirt' in config.keys() else '',
folder_skirt=config['valid_dir_skirt'] if 'valid_dir_skirt' in config.keys() else '',
)
static_dataset_valid = StaticDataset(folder=config['valid_dir'],
@@ -812,6 +829,8 @@ if __name__ == '__main__':
# generate_results()
ckpt = tf.train.Checkpoint(model=model)
manager = tf.train.CheckpointManager(ckpt, 'checkpoints/iccv/' + config['name'], max_to_keep=20)
train_accumulator = Accumulator()