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

small updates, dataset, low dim cvae

Browse Source
This commit is contained in:
luni
2022-09-27 21:21:49 +00:00
parent 83294bc00c
commit a5eff778df
5 changed files with 28 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cvae_low.py
+9 -7
View File
@@ -77,7 +77,7 @@ def train_step(batch):
batch['uv_mask' + res_modifier], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask']], use_mask=True)
pred_enhanced = enhance(upscale_enhancer, pred, batch['uv_mask' + res_modifier], training=False)
pred_full_res = tf.image.resize(pred_enhanced, size=(512, 256), method=tf.image.ResizeMethod.LANCZOS5)
pred_full_res = tf.image.resize(pred_enhanced, size=(512, 256), method=tf.image.ResizeMethod.LANCZOS5) * batch['uv_mask'][:,:,:,tf.newaxis]
mesh_pred = uv2mesh(pred_full_res)
mesh_orig = batch['mesh_unposed']
@@ -87,8 +87,8 @@ def train_step(batch):
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)
e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
c_loss = tf.reduce_mean(collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']))
e_loss = edge_loss(mesh_pred, mesh_orig, batch['mesh_mask'])
c_loss = collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
l1_loss = tf.reduce_sum(tf.reduce_sum((tf.math.abs(batch['uv_unposed' + res_modifier] - pred)), axis=-1), axis=[1, 2])
l2_loss = 1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed' + res_modifier] - pred) ** 2), axis=-1), axis=[1, 2])
l2_loss_full = 1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed'] - pred_full_res) ** 2), axis=-1), axis=[1, 2])
@@ -98,7 +98,7 @@ def train_step(batch):
else:
ms_ssim = tf.image.ssim_multiscale(pred + 1, batch['uv_unposed' + res_modifier] + 1, max_val=3, power_factors=(0.5, 0.3, 0.1, 0.1))
reconstr_loss = l2_loss_full
reconstr_loss = l2_loss_full #+ e_loss + c_loss
reconstr_loss = tf.reduce_mean(reconstr_loss)
kl = kl_loss(mean, logvar)
loss = reconstr_loss + balancing_coeff * kl
@@ -107,6 +107,8 @@ def train_step(batch):
l1_loss = tf.reduce_mean(l1_loss)
l2_loss = tf.reduce_mean(l2_loss)
l2_loss_full = tf.reduce_mean(l2_loss_full)
e_loss = tf.reduce_mean(e_loss)
c_loss = tf.reduce_mean(c_loss)
mean = tf.reduce_mean(mean)
logvar = tf.reduce_mean(logvar)
@@ -215,7 +217,7 @@ def sample(folder):
load_model(model,folder)
count = 0
for batch in dataset_train:
batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic))
batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
nr_of_samples=2
@@ -371,9 +373,9 @@ if __name__ == '__main__':
factory = DatasetFactory()
dataset_train = factory.get_dataset(folder=config['train_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[512, config['resolution']])
dataset_train = factory.get_dataset(folder=config['train_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[512, config['resolution']], in_memory=config['dataset_in_memory'])
static_dataset_train = StaticDataset(folder=config['train_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[512, config['resolution']])
dataset_valid = factory.get_dataset(folder=config['valid_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'])
dataset_valid = factory.get_dataset(folder=config['valid_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[512, config['resolution']], in_memory=config['dataset_in_memory'])
static_dataset_valid = StaticDataset(folder=config['valid_dir'], signature=config['signature_static'])
model = CVAE(config)