add test for ae_static

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')
@@ -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()
+    test()
+    # train()