apply for body and curv

aaa_helpers.txt

@@ -1,4 +1,6 @@
 nohup python auto_encoder_static.py --name="lhunor_ae_static_r7" > nohup/review/ae_static_r7.out 2>&1 &
+nohup python auto_encoder_body.py --name="lhunor_ae_body_r7" > nohup/review/ae_body_r7.out 2>&1 &
+nohup python auto_encoder_curv.py --name="lhunor_ae_curv_r7" > nohup/review/ae_curv_r7.out 2>&1 &
 
 one final convolutional layer for smoothing, on the final coordinates
 

auto_encoder_body.py

@@ -1,6 +1,6 @@
 import os
 import sys
-from configs.config_ae_body_post import config
+from configs.review.config_ae_body import config
 os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
 
 import tensorflow as tf
@@ -34,7 +34,7 @@ from models.cond_encoder import AutoEncoderNew as AutoEncoder
 
 
 def init_wandb():
-    run = wandb.init(project='ae', name=config['name'], notes=config['notes'], dir='/home/workspace/local-data/temp')
+    run = wandb.init(project='ae', name=config['name'], notes=config['notes'], dir='/home/workspace/local-data/temp', tags=config['tags'])
     artifact = wandb.Artifact('code', type='code')
     artifact.add_file('auto_encoder_body.py')
     artifact.add_file('losses.py')
@@ -42,7 +42,7 @@ def init_wandb():
     artifact.add_file('models/cond_encoder.py')
     artifact.add_file('dataset.py')
     artifact.add_file('convnext/models/convnext_tf.py')
-    config_file_names = [key for key in sys.modules.keys()  if key.startswith('configs.')]
+    config_file_names = [key for key in sys.modules.keys()  if key.startswith('configs.review.')]
     for config_file_name in config_file_names:
         artifact.add_file('/'.join(config_file_name.split('.')) + '.py')
     run.log_artifact(artifact)
@@ -137,7 +137,10 @@ def train():
         if step == 0:
             to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
 
-        wandb.log(to_log, step=step)  
+        if step % 1000:    
+            wandb.log(to_log, step=step, commit=True)
+        else:
+            wandb.log(to_log, step=step, commit=False)
             
         if step % 10000 == 0:
             manager_enc.save(checkpoint_number=step)

auto_encoder_curv.py

@@ -1,6 +1,6 @@
 import os
 import sys
-from configs.config_ae_curv_post import config
+from configs.review.config_ae_curv import config
 os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
 
 import tensorflow as tf
@@ -37,7 +37,7 @@ from models.cond_encoder import AutoEncoderNew as AutoEncoder
 testing = False
 def init_wandb():
     if not testing:
-        run = wandb.init(project='ae', name=config['name'], notes=config['notes'], dir='/home/workspace/local-data/temp/')
+        run = wandb.init(project='ae', name=config['name'], notes=config['notes'], dir='/home/workspace/local-data/temp', tags=config['tags'])
         artifact = wandb.Artifact('code', type='code')
         artifact.add_file('auto_encoder_curv.py')
         artifact.add_file('losses.py')
@@ -45,7 +45,7 @@ def init_wandb():
         artifact.add_file('models/cond_encoder.py')
         artifact.add_file('dataset.py')
         artifact.add_file('convnext/models/convnext_tf.py')
-        config_file_names = [key for key in sys.modules.keys()  if key.startswith('configs.')]
+        config_file_names = [key for key in sys.modules.keys()  if key.startswith('configs.review.')]
         for config_file_name in config_file_names:
             artifact.add_file('/'.join(config_file_name.split('.')) + '.py')
         run.log_artifact(artifact)
@@ -178,7 +178,7 @@ def train():
             to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
  
         wandb.log(to_log, step=step)  
-        print(step)
+        # print(step)
 
         if step % 10000 == 0:
             manager_enc.save(checkpoint_number=step)
@@ -206,7 +206,7 @@ def test():
 
 if __name__ == '__main__':
     init_wandb()
-    allow_memory_growth()
+    # allow_memory_growth()
     parser = argparse.ArgumentParser(description='')
     parser.add_argument('--name', type=str)
     args = vars(parser.parse_args())

configs/review/config_ae_body.py

@@ -0,0 +1,51 @@
+config = dict()
+
+
+config['batch_size'] = 4
+config['train_dir'] = '/home/workspace/local-data/train'
+config['valid_dir'] = '/home/workspace/local-data/valid'
+config['test_dir'] = '/home/workspace/data/dataset/test_ready'
+config['use_filled'] = False
+config['resolution'] = 512
+config['is_skirt'] = False
+config['dataset_in_memory'] = False
+config['is_cape'] = False
+config['remove_global_orientation'] = True
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+config['cond_ae'] = dict()
+config['cond_ae']['is_label'] = False
+config['cond_ae']['upscale_output'] = False
+config['cond_ae']['latent_dim'] = config['latent_dim']
+config['cond_ae']['kernel_size'] = 7
+config['cond_ae']['bottleneck_size'] = 16
+config['cond_ae']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae']['encoder_dims'] = [16, 32, 64, 128, config['cond_ae']['bottleneck_size']]
+config['cond_ae']['decoder_depths'] = config['cond_ae']['encoder_depths']
+config['cond_ae']['decoder_dims'] = config['cond_ae']['encoder_dims']
+config['cond_ae']['activation'] = 'tanh'
+
+
+config['gpu'] = '4'
+config['lr'] = 0.0001
+config['name'] = 'body_r7'
+config['notes'] = 'resnet18' 
+config['tags'] = ['review_7']
+
+config['cond_ae']['upscale_output'] = False
+config['pretrained_model'] = ''
+
+
+config['signature_dynamic'] = { 'id': (), 'frame_nr': (),
+                          'pose_shape': (72 + 11,),
+                          'uv_body_posed': (3, 512, 256,),
+                          'mesh_body_posed': (config['smpl_res'], 3,),
+                        }
+config['signature_static'] = { 
+                         'mesh_mask': (config['smpl_res'],),
+                        }

configs/review/config_ae_curv.py

@@ -0,0 +1,50 @@
+config = dict()
+
+
+config['batch_size'] = 4
+config['train_dir'] = '/home/workspace/local-data/train'
+config['valid_dir'] = '/home/workspace/local-data/valid'
+config['test_dir'] = '/home/workspace/data/dataset/test_ready'
+config['use_filled'] = False
+config['resolution'] = 512
+config['is_skirt'] = False
+config['dataset_in_memory'] = False
+config['is_cape'] = False
+config['remove_global_orientation'] = True
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+config['cond_ae'] = dict()
+config['cond_ae']['is_label'] = True
+config['cond_ae']['upscale_output'] = False
+config['cond_ae']['latent_dim'] = config['latent_dim']
+config['cond_ae']['kernel_size'] = 7
+config['cond_ae']['bottleneck_size'] = 40
+config['cond_ae']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae']['encoder_dims'] = [16, 32, 64, 128, config['cond_ae']['bottleneck_size']]
+config['cond_ae']['decoder_depths'] = config['cond_ae']['encoder_depths']
+config['cond_ae']['decoder_dims'] = config['cond_ae']['encoder_dims']
+config['cond_ae']['activation'] = 'tanh'
+
+
+config['gpu'] = '3'
+config['lr'] = 0.001
+config['name'] = 'curv_r7'
+config['notes'] = 'resnet18' 
+config['tags'] = ['review_7']
+
+config['cond_ae']['upscale_output'] = False
+config['pretrained_model'] = ''
+
+
+config['signature_dynamic'] = { 'id': (), 'frame_nr': (),
+                          'uv_curvature': (512, 256, 7),
+                          'mesh_curvature':  (config['smpl_res'])
+                        }
+config['signature_static'] = { 'mesh_mask': (config['smpl_res'],),
+                         'uv_mask': (512, 256,),
+                        }

models/base.py

@@ -152,7 +152,6 @@ class AEEncoderNew(tf.keras.Model):
     def __init__(self, depths, dims, kernel_size=7, activation='tanh'):
         super(AEEncoderNew, self).__init__()
 
-        # self.encoder = ConvNeXt(depths=depths, dims=dims, include_top=False, kernel_size=kernel_size)
         self.encoder = resnet_18()
         self.act = tf.keras.layers.Activation(activation)
 
@@ -180,10 +179,9 @@ class DecoderNew(tf.keras.Model):
         self.is_label = is_label
 
         if not is_label:
-            # self.decoder = ConvNeXtReverse(depths=depths, dims=dims, kernel_size=kernel_size)
             self.decoder = resnet_18(reversed=True)
         else:
-            self.decoder = ConvNeXtReverseLabels(depths=depths, dims=dims, nr_of_classes=7)
+            self.decoder = resnet_18(reversed=True, labels=6)
             
 
     def call(self, inputs, training):

models/resnet.py

@@ -116,9 +116,10 @@ def make_bottleneck_layer(filter_num, blocks, stride=1):
 
 
 class ResNetTypeI(tf.keras.Model):
-    def __init__(self, layer_params, reversed=False):
+    def __init__(self, layer_params, reversed=False, labels=None):
         super(ResNetTypeI, self).__init__()
         self.reversed = reversed
+        self.labels = labels
 
         if not reversed:
             self.conv1 = tf.keras.layers.Conv2D(filters=64,
@@ -130,7 +131,7 @@ class ResNetTypeI(tf.keras.Model):
                                                 strides=2,
                                                 padding="same")
         else:
-            self.conv1 = tf.keras.layers.Conv2DTranspose(filters=3,
+            self.conv1 = tf.keras.layers.Conv2DTranspose(filters=3 if labels is None else labels,
                                                 kernel_size=(7, 7),
                                                 strides=2,
                                                 padding="same")
@@ -177,7 +178,10 @@ class ResNetTypeI(tf.keras.Model):
             x = self.bn1(x, training=training)
             # x = tf.nn.relu(x)
             x = self.conv1(x)
-            x = tf.nn.tanh(x)
+            if self.labels is None:
+                x = tf.nn.tanh(x)
+            # else:
+            #     x = tf.nn.softmax(x)
 
         return x
 
@@ -224,8 +228,8 @@ class ResNetTypeII(tf.keras.Model):
         return output
 
 
-def resnet_18(reversed=False):
+def resnet_18(reversed=False, labels=None):
-    return ResNetTypeI(layer_params=[2, 2, 2, 2], reversed=reversed)
+    return ResNetTypeI(layer_params=[2, 2, 2, 2], reversed=reversed, labels=labels)
 
 
 def resnet_34():