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
auto_encoder_body.py
+3 -9
View File
@@ -21,13 +21,7 @@ from models.cond_encoder import AutoEncoder
def init_wandb():
run = wandb.init(project='vae', name=config['name'],
notes='convnext + convnext, add dense with tanh at the end of reverse convnext,\
flattened bottleneck with 128 latent_dim\ remove global orientation\
fixed gender issue with deoriented bodies, again, aaand again\
trained on 100 garments\
remove multiplications with garment uv mask\
')
run = wandb.init(project='vae', name=config['name'], notes=config['notes'])
artifact = wandb.Artifact('code', type='code')
artifact.add_file('auto_encoder_body.py')
artifact.add_file('losses.py')
@@ -158,9 +152,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=[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=[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=[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=[config['resolution']], in_memory=config['dataset_in_memory'])
static_dataset_valid = StaticDataset(folder=config['valid_dir'], signature=config['signature_static'])
model = AutoEncoder(config['cond_ae'])
cvae.py
+2 -5
View File
@@ -29,7 +29,7 @@ from PIL import Image
from utils import str_to_bool
from tf_utils import check_tf_version, allow_memory_growth, split_gpu_for_testing
from load_models import load_generator, load_discriminator
from losses import d_logistic, d_logistic_r1_reg, g_logistic_non_saturating, g_logistic_ns_pathreg, edge_loss, uv_mask_loss, bend_loss
from losses import edge_loss, uv_mask_loss, bend_loss, kl_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
@@ -39,8 +39,7 @@ from models.cvae_model import CVAE
def init_wandb():
run = wandb.init(project='vae', name=config['name'],
notes=config['notes'])
run = wandb.init(project='vae', name=config['name'], notes=config['notes'])
artifact = wandb.Artifact('code', type='code')
artifact.add_file('cvae.py')
artifact.add_file('models/cvae_model.py')
@@ -322,8 +321,6 @@ if __name__ == '__main__':
parser.add_argument('--name', type=str)
args = vars(parser.parse_args())
batch_size = config['batch_size']
smpl = SMPL()
factory = DatasetFactory()
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)
downscaling_test_curv.py
+1 -1
View File
@@ -104,7 +104,7 @@ if __name__ == '__main__':
'mesh_posed': (14475, 3,),
# 'mesh_unposed': (14475, 3,),
# 'uv_unposed': (3, 512, 256,)
'uv_curvature': (512, 256),
'uv_curvature': (512, 256, 7),
'mesh_curvature': (14475)
}
signature_static = { #'uv_static': (3, 512, 256,),
losses.py
+13
View File
@@ -258,6 +258,19 @@ def bend_loss(x, F, L, mask):
# return loss
def log_normal_pdf(sample, mean, logvar, raxis=1):
log2pi = tf.math.log(2. * np.pi)
res = tf.reduce_sum(-.5 * ((sample - mean) ** 2. * tf.exp(-logvar) + logvar + log2pi), axis=raxis)
return res
def kl_loss(mean, logvar):
kl_loss = 1 + logvar - tf.square(mean) - tf.exp(logvar)
kl_loss = tf.reduce_mean(kl_loss)
kl_loss *= -0.5
return kl_loss
if __name__ == '__main__':
V = tf.ones((4, 14475, 3))
B = tf.zeros((4, 14475, 3))