train curv encoders with 10p and radius 10,25

aaa_helpers.txt

@@ -1,4 +1,6 @@
 nohup python cvae.py --name="lhunor_abl_all_10p_curv_10_r2" > nohup/review/abl_all_10p_curv_10_r2.out 2>&1 &
+nohup python auto_encoder_curv.py --name="lhunor_ae_curv_1_orig_curv_10_10p_r2" > nohup/review/curv_1_orig_curv_10_10p_r2.out 2>&1 &
+nohup python auto_encoder_curv.py --name="lhunor_ae_curv_1_orig_curv_25_10p_r2" > nohup/review/curv_1_orig_curv_25_10p_r2.out 2>&1 &
 
 one final convolutional layer for smoothing, on the final coordinates
 

auto_encoder_curv.py

@@ -1,6 +1,6 @@
 import os
 import sys
-from configs.config_ae_curv_post import config
+from configs.config_ae_curv_1_orig_curv_10_10p_r2 import config
 os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
 
 import tensorflow as tf
@@ -31,13 +31,13 @@ import plotly.express as ex
 from tf_utils import check_tf_version, allow_memory_growth
 from dataset import DatasetFactory, StaticDataset
 from custom_utils.custom_utils import batch_to_dict
-from models.cond_encoder import AutoEncoderNew as AutoEncoder
+from models.cond_encoder import 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')
@@ -47,7 +47,9 @@ def init_wandb():
         artifact.add_file('convnext/models/convnext_tf.py')
         config_file_names = [key for key in sys.modules.keys()  if key.startswith('configs.')]
         for config_file_name in config_file_names:
-            artifact.add_file('/'.join(config_file_name.split('.')) + '.py')
+            file_name = '/'.join(config_file_name.split('.')) + '.py'
+            if os.path.exists(file_name):
+                artifact.add_file(file_name)
         run.log_artifact(artifact)
 
 uv2mesh = UV2MeshLabel(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
@@ -99,43 +101,70 @@ def dice_dense(prediction, ground_truth, weight_map=None):
 def train_step(batch):	
     with tf.GradientTape() as tape:	
         pred = model([batch['uv_curvature'], batch['uv_mask']], training=True)	
-
         ce_uv_loss = scce(tf.math.argmax(batch['uv_curvature'], axis=-1), pred)	
         ms_ssim = tf.reduce_mean(tf.image.ssim(pred, batch['uv_curvature'], max_val=1))	
-
+        loss = ce_uv_loss	
-        mesh_pred = uv2mesh(pred) * batch['mesh_mask'][:,:,tf.newaxis]
+    pred_mesh = tf.math.argmax(uv2mesh(pred), axis=-1)	
-        mesh_orig = tf.cast(uv2mesh(tf.math.argmax(batch['uv_curvature'], axis=-1, output_type=tf.int32)), tf.float32) * batch['mesh_mask']
-        ce_mesh_loss = scce(mesh_orig, mesh_pred)
-
-        loss = ce_uv_loss + ce_mesh_loss
-
     mean_iou = get_mean_iou(tf.math.argmax(batch['uv_curvature'], axis=-1), tf.math.argmax(pred, axis=-1))	
+    mean_iou_mesh = get_mean_iou(tf.cast(batch['mesh_curvature'], tf.float32) * batch['mesh_mask'], tf.cast(pred_mesh, tf.float32) * batch['mesh_mask'])	
     dice = dice_dense(pred, batch['uv_curvature'])	
     	
     gradients = tape.gradient(loss, model.trainable_variables)	
     optimizer.apply_gradients(zip(gradients, model.trainable_variables))	
-    return { 'train/loss': loss, 'train/ce_uv_loss': ce_uv_loss, 'train/ssim': ms_ssim, 'train/mean_iou': mean_iou,
+    return { 'train/loss': loss, 'train/ce_uv_loss': ce_uv_loss, 'train/ssim': ms_ssim, 'train/mean_iou': mean_iou, 'train/mean_iou_mesh': mean_iou_mesh, 'train/dice': dice }	
-             'train/dice': dice, 'train/ce_mesh_loss': ce_mesh_loss }
-        
-
         	
 @tf.function	
 def valid_step(batch):	
     pred = model([batch['uv_curvature'], batch['uv_mask']], training=False)	
-
     ce_uv_loss = scce(tf.math.argmax(batch['uv_curvature'], axis=-1), pred)	
     ms_ssim = tf.reduce_mean(tf.image.ssim(pred, batch['uv_curvature'], max_val=1))	
-
+    loss = ce_uv_loss	
-    mesh_pred = uv2mesh(pred) * batch['mesh_mask'][:,:,tf.newaxis]
+    pred_mesh = tf.math.argmax(uv2mesh(pred), axis=-1)	
-    mesh_orig = tf.cast(uv2mesh(tf.math.argmax(batch['uv_curvature'], axis=-1, output_type=tf.int32)), tf.float32) * batch['mesh_mask']
-    ce_mesh_loss = scce(mesh_orig, mesh_pred)
-    loss = ce_uv_loss + ce_mesh_loss
-
     mean_iou = get_mean_iou(tf.math.argmax(batch['uv_curvature'], axis=-1), tf.math.argmax(pred, axis=-1))	
+    mean_iou_mesh = get_mean_iou(tf.cast(batch['mesh_curvature'], tf.float32) * batch['mesh_mask'], tf.cast(pred_mesh, tf.float32) * batch['mesh_mask'])	
     dice = dice_dense(pred, batch['uv_curvature'])	
+    return { 'valid/loss': loss, 'valid/ce_uv_loss': ce_uv_loss, 'valid/ssim': ms_ssim, 'valid/mean_iou': mean_iou, 'valid/mean_iou_mesh': mean_iou_mesh, 'valid/dice': dice }
 
-    return { 'valid/loss': loss, 'valid/ce_uv_loss': ce_uv_loss, 'valid/ssim': ms_ssim, 'valid/mean_iou': mean_iou,
+# def train_step(batch):
-             'valid/dice': dice, 'valid/ce_mesh_loss': ce_mesh_loss }
+#     with tf.GradientTape() as tape:
+#         pred = model([batch['uv_curvature'], batch['uv_mask']], training=True)
+
+#         ce_uv_loss = scce(tf.math.argmax(batch['uv_curvature'], axis=-1), pred)
+#         ms_ssim = tf.reduce_mean(tf.image.ssim(pred, batch['uv_curvature'], max_val=1))
+
+#         mesh_pred = uv2mesh(pred) * batch['mesh_mask'][:,:,tf.newaxis]
+#         mesh_orig = tf.cast(uv2mesh(tf.math.argmax(batch['uv_curvature'], axis=-1, output_type=tf.int32)), tf.float32) * batch['mesh_mask']
+#         ce_mesh_loss = scce(mesh_orig, mesh_pred)
+
+#         loss = ce_uv_loss + ce_mesh_loss
+
+#     mean_iou = get_mean_iou(tf.math.argmax(batch['uv_curvature'], axis=-1), tf.math.argmax(pred, axis=-1))
+#     dice = dice_dense(pred, batch['uv_curvature'])
+    
+#     gradients = tape.gradient(loss, model.trainable_variables)
+#     optimizer.apply_gradients(zip(gradients, model.trainable_variables))
+#     return { 'train/loss': loss, 'train/ce_uv_loss': ce_uv_loss, 'train/ssim': ms_ssim, 'train/mean_iou': mean_iou,
+#              'train/dice': dice, 'train/ce_mesh_loss': ce_mesh_loss }
+        
+
+
+# @tf.function
+# def valid_step(batch):
+#     pred = model([batch['uv_curvature'], batch['uv_mask']], training=False)
+
+#     ce_uv_loss = scce(tf.math.argmax(batch['uv_curvature'], axis=-1), pred)
+#     ms_ssim = tf.reduce_mean(tf.image.ssim(pred, batch['uv_curvature'], max_val=1))
+
+#     mesh_pred = uv2mesh(pred) * batch['mesh_mask'][:,:,tf.newaxis]
+#     mesh_orig = tf.cast(uv2mesh(tf.math.argmax(batch['uv_curvature'], axis=-1, output_type=tf.int32)), tf.float32) * batch['mesh_mask']
+#     ce_mesh_loss = scce(mesh_orig, mesh_pred)
+#     loss = ce_uv_loss + ce_mesh_loss
+
+#     mean_iou = get_mean_iou(tf.math.argmax(batch['uv_curvature'], axis=-1), tf.math.argmax(pred, axis=-1))
+#     dice = dice_dense(pred, batch['uv_curvature'])
+
+#     return { 'valid/loss': loss, 'valid/ce_uv_loss': ce_uv_loss, 'valid/ssim': ms_ssim, 'valid/mean_iou': mean_iou,
+#              'valid/dice': dice, 'valid/ce_mesh_loss': ce_mesh_loss }
         
 
 def plot_everything(batch: dict, type: str):
@@ -206,7 +235,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())
@@ -226,7 +255,8 @@ if __name__ == '__main__':
                                         in_memory=config['dataset_in_memory'],
                                         use_no_orient=config['remove_global_orientation'],
                                         frames_file=config['frames_file'] if 'frames_file' in config.keys() else None, 
-                                        is_cape=config['is_cape'])
+                                        is_cape=config['is_cape'],
+                                        curv_neighbourhood=config['curv_neighbourhood'])
     static_dataset_train = StaticDataset(folder=config['train_dir'],
                                          signature=config['signature_static'], 
                                          resolutions=[config['resolution']], 
@@ -240,7 +270,8 @@ if __name__ == '__main__':
                                         in_memory=config['dataset_in_memory'], 
                                         use_no_orient=config['remove_global_orientation'], 
                                         frames_file=config['frames_file_valid'] if 'frames_file_valid' in config.keys() else None, 
-                                        is_cape=config['is_cape'])
+                                        is_cape=config['is_cape'],
+                                        curv_neighbourhood=config['curv_neighbourhood'])
     static_dataset_valid = StaticDataset(folder=config['valid_dir'],
                                          signature=config['signature_static'], 
                                          resolutions=[config['resolution']], 

configs/config_ae_curv.py

@@ -1,13 +1,23 @@
 from configs.config_base import config
 
-config['gpu'] = '5'
+config['gpu'] = '4'
 config['lr'] = 0.001
-config['name'] = 'curv_1' #'ae_curv_9'
+config['name'] = 'curv_1_orig_r2' #'ae_curv_9'
 config['notes'] = 'convnext + convnext, 16*8*32, ce loss. flattened latent dim of 128. trained on 100 garments. switched back to sparse CCE, normal one was giving NaN | \
 | \
   '
+config['tags'] = ['review_2']
+config['remove_global_orientation'] = True
 
 config['cond_ae']['is_label'] = True
+config['is_skirt'] = True
+
+config['train_dir'] = '/home/workspace/data/dataset/local-cache/train_skirt'
+config['valid_dir'] = '/home/workspace/local-data/valid_skirt'
+config['test_dir'] = '/home/workspace/data/dataset/test_skirt_ready_2'
+config['frames_file'] = 'frames_skirt_10p.txt'
+config['frames_file_valid'] = 'frames_skirt_10p.txt'
+config['curv_neighbourhood'] = 10
 
 config['signature_dynamic'] = { 'id': (), 'frame_nr': (),
                           'uv_curvature': (512, 256, 7),

configs/config_ae_curv_1_orig_curv_10_10p_r2.py

@@ -0,0 +1,30 @@
+from configs.config_base import config
+
+config['gpu'] = '4'
+config['lr'] = 0.001
+config['name'] = 'curv_1_orig_curv_10_10p_r2' #'ae_curv_9'
+config['notes'] = 'convnext + convnext, 16*8*32, ce loss. flattened latent dim of 128. trained on 100 garments. switched back to sparse CCE, normal one was giving NaN | \
+| \
+  '
+config['tags'] = ['review_2']
+config['remove_global_orientation'] = True
+config['dataset_in_memory'] = True
+
+config['cond_ae']['is_label'] = True
+config['is_skirt'] = True
+
+config['train_dir'] = '/home/workspace/data/dataset/local-cache/train_skirt'
+config['valid_dir'] = '/home/workspace/local-data/valid_skirt'
+config['test_dir'] = '/home/workspace/data/dataset/test_skirt_ready_2'
+config['frames_file'] = 'frames_skirt_10p.txt'
+config['frames_file_valid'] = 'frames_skirt_10p.txt'
+config['curv_neighbourhood'] = 10
+
+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,),
+                        }
+

configs/config_ae_curv_1_orig_curv_25_10p_r2.py

@@ -0,0 +1,30 @@
+from configs.config_base import config
+
+config['gpu'] = '6'
+config['lr'] = 0.001
+config['name'] = 'curv_1_orig_curv_25_10p_r2' #'ae_curv_9'
+config['notes'] = 'convnext + convnext, 16*8*32, ce loss. flattened latent dim of 128. trained on 100 garments. switched back to sparse CCE, normal one was giving NaN | \
+| \
+  '
+config['tags'] = ['review_2']
+config['remove_global_orientation'] = True
+config['dataset_in_memory'] = True
+
+config['cond_ae']['is_label'] = True
+config['is_skirt'] = True
+
+config['train_dir'] = '/home/workspace/data/dataset/local-cache/train_skirt'
+config['valid_dir'] = '/home/workspace/local-data/valid_skirt'
+config['test_dir'] = '/home/workspace/data/dataset/test_skirt_ready_2'
+config['frames_file'] = 'frames_skirt_10p.txt'
+config['frames_file_valid'] = 'frames_skirt_10p.txt'
+config['curv_neighbourhood'] = 25
+
+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,),
+                        }
+

convnext/models/convnext_tf.py

@@ -366,8 +366,8 @@ class ConvNeXtReverseLabels(tf.keras.Model):
             self.stages.append(stage)
             cur += depths[i]
 
-        # self.last = Dense(nr_of_classes, activation='softmax')
+        self.last = Dense(nr_of_classes, activation='softmax')
-        self.last = Dense(nr_of_classes - 1)
+        # self.last = Dense(nr_of_classes - 1)
 
     def call(self, x, training):
         for i in reversed(range(len(self.downsample_layers))):

dataset.py

@@ -261,7 +261,7 @@ class DatasetFactory():
 
 
             if 'mesh_curvature' in signature.keys():
-                mesh_curvature = np.load(os.path.join(folder, f'mesh_curvature_{curv_neighbourhood}', id + separator + frame_nr + '.npz')); mesh_curvature = mesh_curvature['mesh_curvature']
+                mesh_curvature = np.load(os.path.join(folder, f'mesh_curvature_{curv_neighbourhood}' if curv_neighbourhood != 25 else 'mesh_curvature', id + separator + frame_nr + '.npz')); mesh_curvature = mesh_curvature['mesh_curvature']
                 item['mesh_curvature'] = mesh_curvature.astype(np.int32)
 
 
@@ -271,7 +271,8 @@ class DatasetFactory():
                     res_string = '_' + str(res)
 
                 if 'uv_curvature' + res_string in signature.keys():
-                    uv_curvature = np.load(os.path.join(folder, f'uv_curvature_{curv_neighbourhood}' + res_string, id + separator + frame_nr + '.npz')); 
+                    print(os.path.join(folder, (f'uv_curvature_{curv_neighbourhood}' if curv_neighbourhood != 25 else 'uv_curvature')+ res_string, id + separator + frame_nr + '.npz'))
+                    uv_curvature = np.load(os.path.join(folder, (f'uv_curvature_{curv_neighbourhood}' if curv_neighbourhood != 25 else 'uv_curvature')+ res_string, id + separator + frame_nr + '.npz')); 
                     if res_string != '':
                         uv_curvature = uv_curvature['uv_curvature'] 
                     else:
@@ -382,15 +383,16 @@ if __name__ == '__main__':
                         #  'uv_shaped_body': (3, 512, 256,),
                           }
 
-    dataset = factory.get_dataset(folder='/home/workspace/local-data/train_skirt',
+    dataset = factory.get_dataset(folder='/home/workspace/data/dataset/local-cache/train_skirt',
                                   batch_size=1,
                                   signature=signature_dynamic,
                                   resolutions=[512],
                                   do_repeat=False,
                                   use_one_hot=False,
                                   is_cape=False,
-                                  frames_file='frames_dress_30k.txt')
+                                  frames_file='frames_skirt_10p.txt',
-    static_dataset = StaticDataset(folder='/home/workspace/local-data/train_skirt',
+                                  curv_neighbourhood=25)
+    static_dataset = StaticDataset(folder='/home/workspace/data/dataset/local-cache/train_skirt',
                                    signature=signature_static,
                                    resolutions=[512],
                                   is_cape=False)