train curv encoders with 10p and radius 10,25

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'])
@@ -96,46 +98,73 @@ def dice_dense(prediction, ground_truth, weight_map=None):
 
 
 @tf.function
-def train_step(batch):
-    with tf.GradientTape() as tape:
-        pred = model([batch['uv_curvature'], batch['uv_mask']], training=True)
+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	
+    pred_mesh = tf.math.argmax(uv2mesh(pred), axis=-1)	
+    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, 'train/mean_iou_mesh': mean_iou_mesh, 'train/dice': dice }	
+        	
+@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	
+    pred_mesh = tf.math.argmax(uv2mesh(pred), axis=-1)	
+    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 }
 
-        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))
+# def train_step(batch):
+#     with tf.GradientTape() as tape:
+#         pred = model([batch['uv_curvature'], batch['uv_mask']], training=True)
 
-        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)
+#         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 + ce_mesh_loss
+#         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)
 
-    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'])
+#         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 }
+#     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)
+# @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))
+#     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
+#     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'])
+#     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 }
+#     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']],