add testing to static ae

auto_encoder_static.py

@@ -8,7 +8,7 @@ import wandb
 import random
 import numpy as np
 import tensorflow as tf
-from custom_utils.custom_utils import UV2Mesh, maskFacesBatch
+from custom_utils.custom_utils import UV2Mesh, load_model, maskFacesBatch
 import plotly.graph_objects as go
 import plotly.express as ex
 
@@ -43,7 +43,7 @@ uv2mesh = UV2Mesh(res=256)
 def train_step(batch):
 
     with tf.GradientTape() as tape:
-        pred = model([batch['uv_static'], batch['uv_mask']])
+        pred = model([batch['uv_static'], batch['uv_mask']], training=True)
 
         mesh_pred = uv2mesh(pred)
         mesh_orig = batch['mesh_static']
@@ -134,8 +134,24 @@ def train():
 
         step += 1
 
+
+def test():
+    load_model(model.encoder, 'checkpoints/' + config['name'] + '_enc', assert_consumed_on=True)
+    load_model(model.decoder, 'checkpoints/' + config['name'] + '_dec', assert_consumed_on=True)
+    for batch in dataset_train:
+        batch = static_dataset_train.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()
+
+        ex.imshow((img[0] + 1) / 2).show()
+        print()
+
+
 if __name__ == '__main__':
-    init_wandb()
+    # init_wandb()
     allow_memory_growth()
     parser = argparse.ArgumentParser(description='')
     parser.add_argument('--name', type=str)
@@ -162,4 +178,5 @@ if __name__ == '__main__':
     ckpt_dec = tf.train.Checkpoint(model=model.decoder)
     manager_dec = tf.train.CheckpointManager(ckpt_dec, 'checkpoints/' + config['name'] + '_dec', max_to_keep=3)
 
-    train()
+    test()
+    # train()