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

commit all

Browse Source
This commit is contained in:
luni
2023-01-25 12:25:59 +00:00
parent 1bf224841f
commit 7470ca9fb6
139 changed files with 82069 additions and 811 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cvae_pyramid.py
+132 -26
View File
@@ -1,6 +1,6 @@
import os
import sys
from configs.config_abl_pyr_64_128_cape_m import config
from configs.post.config_abl_pyr_64_128_5 import config
os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
import tensorflow as tf
@@ -35,18 +35,18 @@ if config['is_skirt']:
else:
from losses import edge_loss
from losses import uv_mask_loss, bend_loss, kl_loss, collision_ratio
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 dataset import DatasetFactory, StaticDataset
from custom_utils.smpl import SMPL
from models.pyramid_model_optimized import Pyramid
from models.pyramid_model_optimized import PyramidNew as Pyramid
from postprocessing import fix_collision
def init_wandb():
run = wandb.init(project='ablation', name=config['name'], notes=config['notes'], dir='/wandb')
run = wandb.init(project='ablation', name=config['name'], notes=config['notes'], dir='/wandb', tags=config['tags'])
artifact = wandb.Artifact('code', type='code')
artifact.add_file('cvae_pyramid.py')
@@ -59,7 +59,7 @@ def init_wandb():
artifact.add_file('models/cond_encoder.py')
artifact.add_file('dataset.py')
artifact.add_file('convnext/models/convnext_tf.py')
config_file_names = [key for key in sys.modules.keys() if key.startswith('configs.')]
config_file_names = [key for key in sys.modules.keys() if key.startswith('configs.post.')]
for config_file_name in config_file_names:
artifact.add_file('/'.join(config_file_name.split('.')) + '.py')
run.log_artifact(artifact)
@@ -208,7 +208,6 @@ def plot_3d(mesh_gt, mesh_pred_unposed_levels, faces, mesh_posed, collision, col
return fig
@tf.function
def train_step(batch):
levels_input, conds = group_inputs_for_levels(batch, config)
@@ -251,7 +250,8 @@ def train_step(batch):
l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((uv_orig - pred) ** 2), axis=-1), axis=[1, 2]))
# ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 5, batch['uv_unposed'] + 5, max_val=10, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
# ms_ssim = tf.reduce_mean(tf.image.ssim(pred + 1, batch['uv_unposed'] + 1, max_val=3))
l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1), axis=-1))
l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, axis=-1), axis=-1))
n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_mask'], is_skirt=config['is_skirt'])
kl = kl_loss(mean, logvar)
if config['use_extra_losses']:
@@ -261,6 +261,8 @@ def train_step(batch):
else:
reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (l2_loss)
if config['use_norm_loss'] and step > 2000:
reconstr_loss_total = reconstr_loss_total + n_loss * config['reg_loss_coeff']
kl_total += kl
@@ -271,7 +273,7 @@ def train_step(batch):
output = { **output, **{ 'train/edge_loss' + res_modifier: e_loss, 'train/collision_loss' + res_modifier: c_loss,'train/euclid_error' + res_modifier: metric,
'train/mean' + res_modifier: mean, 'train/logvar' + res_modifier: logvar, 'train/kl_loss' + res_modifier: kl,
'train/var' + res_modifier: np.e ** logvar, 'train/l2_loss' + res_modifier: l2_loss,
'train/var' + res_modifier: np.e ** logvar, 'train/l2_loss' + res_modifier: l2_loss, 'train/normal_loss': n_loss,
'train/euclid_error': metric, 'train/l2_mesh_loss': l2_mesh_loss } } # 'train/ms_ssim' + res_modifier: ms_ssim,
loss = reconstr_loss_total + balancing_coeff * kl_total
@@ -321,7 +323,8 @@ def valid_step(batch):
l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((uv_orig - pred) ** 2), axis=-1), axis=[1, 2]))
# ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 5, batch['uv_unposed'] + 5, max_val=10, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
# ms_ssim = tf.reduce_mean(tf.image.ssim(pred + 1, batch['uv_unposed'] + 1, max_val=3))
l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1), axis=-1))
l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, axis=-1), axis=-1))
n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_mask'], is_skirt=config['is_skirt'])
kl = kl_loss(mean, logvar)
if config['use_extra_losses']:
@@ -331,6 +334,9 @@ def valid_step(batch):
else:
reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (l2_loss)
if config['use_norm_loss'] and step > 2000:
reconstr_loss_total = reconstr_loss_total + n_loss * config['reg_loss_coeff']
kl_total += kl
mean = tf.reduce_mean(mean)
@@ -340,7 +346,7 @@ def valid_step(batch):
output = { **output, **{ 'valid/edge_loss' + res_modifier: e_loss, 'valid/collision_loss' + res_modifier: c_loss,'valid/euclid_error' + res_modifier: metric,
'valid/mean' + res_modifier: mean, 'valid/logvar' + res_modifier: logvar, 'valid/kl_loss' + res_modifier: kl,
'valid/var' + res_modifier: np.e ** logvar, 'valid/l2_loss' + res_modifier: l2_loss,
'valid/var' + res_modifier: np.e ** logvar, 'valid/l2_loss' + res_modifier: l2_loss, 'valid/normal_loss': n_loss,
'valid/euclid_error': metric, 'valid/l2_mesh_loss': l2_mesh_loss } } # 'valid/ms_ssim' + res_modifier: ms_ssim,
loss = reconstr_loss_total + balancing_coeff * kl_total
@@ -669,13 +675,73 @@ def test():
pass
def generate_results():
load_model(model, 'checkpoints/' + config['name'])
dataset_test = factory.get_dataset(folder=config['test_dir'],
batch_size=config['batch_size'],
# batch_size=1,
signature=config['signature_dynamic'],
resolutions=dataset_resolutions,
in_memory=config['dataset_in_memory'],
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'],
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'])
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']))
levels_input, conds = group_inputs_for_levels(batch, config)
out_levels = model([levels_input, conds], training=False)
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])])
i = -1
pred = out_levels[i][0]
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)
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)
pass
@tf.function
def speed_step(levels_input, conds):
return model.sample([levels_input, conds], training=False)
if __name__ == '__main__':
init_wandb()
parser = argparse.ArgumentParser(description='')
parser.add_argument('--name', type=str)
parser.add_argument('--name', type=str, default='lhunor')
args = vars(parser.parse_args())
# config_file_names = [key for key in sys.modules.keys() if key.startswith('configs.')]
# assert len(config_file_names) == 1
# config_file_name = config_file_names[0]
# # assert config_file_name.split('.')[1].endswith(args.name)
# assert config_file_name.split('.')[1].endswith(config['name'])
smpl = SMPL(orig_smpl=config['is_cape'])
skirt_modifier = '_skirt' if config['is_skirt'] else ''
dataset_resolutions = config['pyramid']['resolutions'].copy()
@@ -683,10 +749,31 @@ if __name__ == '__main__':
dataset_resolutions.append(512)
factory = DatasetFactory()
dataset_train = factory.get_dataset(folder=config['train_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], resolutions=dataset_resolutions, in_memory=config['dataset_in_memory'], is_cape=config['is_cape'])
static_dataset_train = StaticDataset(folder=config['train_dir'], signature=config['signature_static'], resolutions=dataset_resolutions, is_cape=config['is_cape'])
dataset_valid = factory.get_dataset(folder=config['valid_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], resolutions=dataset_resolutions, in_memory=config['dataset_in_memory'], is_cape=config['is_cape'])
static_dataset_valid = StaticDataset(folder=config['valid_dir'], signature=config['signature_static'], resolutions=dataset_resolutions, is_cape=config['is_cape'])
dataset_train = factory.get_dataset(folder=config['train_dir'],
batch_size=config['batch_size'],
# batch_size=1,
signature=config['signature_dynamic'],
resolutions=dataset_resolutions,
in_memory=config['dataset_in_memory'],
is_cape=config['is_cape'],
frames_file=config['frames_file'] if 'frames_file' in config.keys() else None)
static_dataset_train = StaticDataset(folder=config['train_dir'],
signature=config['signature_static'],
resolutions=dataset_resolutions,
in_memory=True,
is_cape=config['is_cape'])
dataset_valid = factory.get_dataset(folder=config['valid_dir'],
batch_size=config['batch_size'],
signature=config['signature_dynamic'],
resolutions=dataset_resolutions,
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)
static_dataset_valid = StaticDataset(folder=config['valid_dir'],
signature=config['signature_static'],
resolutions=dataset_resolutions,
in_memory=True,
is_cape=config['is_cape'])
model: Pyramid = Pyramid(config)
@@ -700,34 +787,38 @@ if __name__ == '__main__':
# sample('checkpoints/cvae_13')
# test()
# generate_results()
ckpt = tf.train.Checkpoint(model=model)
manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + config['name'], max_to_keep=20)
manager = tf.train.CheckpointManager(ckpt, 'checkpoints/post/' + config['name'], max_to_keep=20)
# load_model(model, 'checkpoints/pyr_abl_32_skirt', iteration='180000')
if config['pretrained_model'] != '':
# assert len(config['pyramid']['resolutions']) == 1
load_model(model, config['pretrained_model'], iteration=config['pretrained_model_iteration'])
if len(config['pyramid']['resolutions']) == 1:
model.static_enc.trainable = False
model.body_enc.trainable = False
model.curv_enc.trainable = False
else:
load_model(model, config['pyramid']['pretrained_pyramid'])
load_model(model, config['pyramid']['pretrained_pyramid'], iteration=(config['pyramid']['pretrained_pyramid_iteration'] if 'pretrained_pyramid_iteration' in config.keys() else ''))
model.static_enc.trainable = False
model.body_enc.trainable = False
model.curv_enc.trainable = False
model.cond_encoder.trainable = False
# load_model(model.levels[0].upscale_enhancer, config['pyramid']['levels'][0]['upscale_enhancer'])
# model.levels[0].upscale_enhancer.trainable = False
# the last level doesn't have an upscale enhancer as it's enhanced for the final resolutin directly
# might not be necessary, as enhancing to final or a bit higher resolution seems to be the same
for i in range(len(model.levels)-1):
for i in range(len(model.levels)):
load_model(model.levels[i].upscale_enhancer, config['pyramid']['levels'][i]['upscale_enhancer'])
model.levels[i].upscale_enhancer.trainable = False
if config['pyramid']['resolutions'][-1] != 512:
load_model(model.upscale_enhancer, config['pyramid']['final_upscale_enhancer'])
model.upscale_enhancer.trainable = False
# if config['pyramid']['resolutions'][-1] != 512:
# load_model(model.upscale_enhancer, config['pyramid']['final_upscale_enhancer'])
# model.upscale_enhancer.trainable = False
for i in range(len(model.levels)-1):
model.levels[i].trainable = False
@@ -738,9 +829,22 @@ if __name__ == '__main__':
step = 0
times = []
for batch in dataset_train:
batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
# levels_input, conds = group_inputs_for_levels(batch, config)
# start = time.perf_counter()
# out = speed_step(levels_input, conds)
# t = time.perf_counter() - start
# print(t)
# times.append(t)
# if step > 200:
# break
train_loss = train_step(batch)
data_valid = valid_iterator.get_next()
@@ -753,7 +857,7 @@ if __name__ == '__main__':
valid_figs = plot_everything(data_valid, 'valid')
to_log = { **to_log, **train_figs, **valid_figs }
if step % 10000 == 0:
if step % 5000 == 0:
manager.save(checkpoint_number=step)
if step == 0:
to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
@@ -770,4 +874,6 @@ if __name__ == '__main__':
print(step)
step += 1
step += 1
# print('mean:')
# print(np.mean(times[1:]))