From 3dac84e189a06a301dd1ac6ce13a92333e05ce1d Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Wed, 12 Oct 2022 21:07:49 +0000 Subject: [PATCH] add test for ae_static --- auto_encoder_static.py | 53 ++++++++++++++++++++++++++++-------------- 1 file changed, 35 insertions(+), 18 deletions(-) diff --git a/auto_encoder_static.py b/auto_encoder_static.py index e8feba6..9509d8f 100644 --- a/auto_encoder_static.py +++ b/auto_encoder_static.py @@ -38,9 +38,7 @@ from models.cond_encoder import AutoEncoder def init_wandb(): run = wandb.init(project='ae', name=config['name'], - notes='convnext + convnext, add dense with tanh at the end of reverse convnext,\ - flattened latent dim to 128\ - train on 100 garments') + notes=config['notes']) artifact = wandb.Artifact('code', type='code') artifact.add_file('auto_encoder_static.py') artifact.add_file('losses.py') @@ -143,7 +141,7 @@ def train(): wandb.log(to_log, step=step) - + if step % 10000 == 0: manager_enc.save(checkpoint_number=step) manager_dec.save(checkpoint_number=step) @@ -156,20 +154,39 @@ def train(): def test(): load_model(model.encoder, 'checkpoints/' + config['name'] + '_enc') load_model(model.decoder, 'checkpoints/' + config['name'] + '_dec') - for batch in dataset_train: - batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic'])) + + dataset_test = factory.get_dataset(folder=config['test_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=False, do_repeat=False) + static_dataset_test = StaticDataset(folder=config['test_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=False) + + avg_error = 0 + count = 0 + + for batch in dataset_test: + batch = static_dataset_test.expand_batch(batch_to_dict(batch, config['signature_dynamic'])) pred = model([batch['uv_static'], batch['uv_mask']], training=False) - orig_img = batch['uv_static'] - pred_img = pred - img = tf.concat((orig_img, pred_img), axis=2).numpy() + # orig_img = batch['uv_static'] + # pred_img = pred + # img = tf.concat((orig_img, pred_img), axis=2).numpy() - ex.imshow((img[0] + 1) / 2).show() - print() + # ex.imshow((img[0] + 1) / 2).show() + + + + mesh_pred = uv2mesh(pred) + mesh_orig = batch['mesh_static'] + metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask']) + + avg_error += metric + count += 1 + + print(count) + + print() if __name__ == '__main__': - init_wandb() + # init_wandb() allow_memory_growth() parser = argparse.ArgumentParser(description='') parser.add_argument('--name', type=str) @@ -192,10 +209,10 @@ if __name__ == '__main__': model : AutoEncoder = AutoEncoder(config['cond_ae']) optimizer = tf.optimizers.Adam(learning_rate=config['lr']) - ckpt_enc = tf.train.Checkpoint(model=model.encoder) - manager_enc = tf.train.CheckpointManager(ckpt_enc, 'checkpoints/' + config['name'] + '_enc', max_to_keep=3) - ckpt_dec = tf.train.Checkpoint(model=model.decoder) - manager_dec = tf.train.CheckpointManager(ckpt_dec, 'checkpoints/' + config['name'] + '_dec', max_to_keep=3) + # ckpt_enc = tf.train.Checkpoint(model=model.encoder) + # manager_enc = tf.train.CheckpointManager(ckpt_enc, 'checkpoints/' + config['name'] + '_enc', max_to_keep=3) + # ckpt_dec = tf.train.Checkpoint(model=model.decoder) + # manager_dec = tf.train.CheckpointManager(ckpt_dec, 'checkpoints/' + config['name'] + '_dec', max_to_keep=3) - # test() - train() \ No newline at end of file + test() + # train() \ No newline at end of file