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()

configs/config_ae_static.py

@@ -2,7 +2,7 @@ from configs.config_base import config
 
 config['gpu'] = '0'
 config['lr'] = 0.0005
-config['name'] = 'ae_static_7'
+config['name'] = 'ae_static_7_x'
 
 config['signature_dynamic'] = { 'id': (), }
 config['signature_static'] = { 'uv_static': (512, 256, 3),

custom_utils/custom_utils.py

@@ -203,10 +203,13 @@ def sparse_to_tensor(L):
     return tf.SparseTensor(idx, data, shape)
 
 
-def load_model(model, checkpoint, disable_assert=False):
+def load_model(model, checkpoint, disable_assert=False, assert_consumed_on=False):
     ckpt = tf.train.Checkpoint(model=model)
     manager = tf.train.CheckpointManager(ckpt, checkpoint, max_to_keep=1)
     if not disable_assert:
         assert manager.latest_checkpoint is not None
         assert manager.latest_checkpoint.split('/')[-1] != 'ckpt-0'
+    if assert_consumed_on:
+        ckpt.restore(manager.latest_checkpoint).assert_consumed()
+    else:
         ckpt.restore(manager.latest_checkpoint).assert_existing_objects_matched()#.assert_consumed()

models/base.py

@@ -1,7 +1,4 @@
-import os
 import tensorflow as tf
-
-from custom_utils.custom_utils import load_model
 from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block, ConvNeXtReverseLabels
 
 class Bottleneck(tf.keras.Model):
@@ -13,7 +10,6 @@ class Bottleneck(tf.keras.Model):
                 [Block(dim=dim, drop_path=0, kernel_size=kernel_size,
                        layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
             )
-        # self.bottleneck = [Block(dim=dim, drop_path=0, kernel_size=kernel_size, layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
             
         self.flatten = tf.keras.layers.Flatten()
         self.dense = tf.keras.layers.Dense(latent_dim)
@@ -21,9 +17,7 @@ class Bottleneck(tf.keras.Model):
     def call(self, inputs, training=None):
         x = inputs
 
-        # for bt in self.bottleneck:
+        x = self.bottleneck(x, training=training)
-        #     x = bt(x, training=training)
-        x = self.bottleneck(x)
         x = self.flatten(x)
         x = self.dense(x)
 
@@ -32,7 +26,7 @@ class Bottleneck(tf.keras.Model):
 
 class VariationalEncoder(tf.keras.Model):
     
-    def __init__(self, latent_dim, depths, dims, use_bottleneck=True, kernel_size=7):
+    def __init__(self, latent_dim, depths, dims, kernel_size=7):
         super(VariationalEncoder, self).__init__()
 
         self.encoder = ConvNeXt(depths=depths, dims=dims, include_top=False, kernel_size=kernel_size)