update ae curv

auto_encoder_curv.py

@@ -1,15 +1,28 @@
 import os
 from configs.config_ae_curv import config
 os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
+
+import tensorflow as tf
+gpus = tf.config.experimental.list_physical_devices('GPU')
+if gpus:
+    try:
+        # Currently, memory growth needs to be the same across GPUs
+        for gpu in gpus:
+            tf.config.experimental.set_memory_growth(gpu, True)
+        logical_gpus = tf.config.experimental.list_logical_devices('GPU')
+        print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
+    except RuntimeError as e:
+        # Memory growth must be set before GPUs have been initialized
+        print(e)
+
 import time
 import argparse
 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, UV2MeshLabel, maskFacesBatch
 import plotly.graph_objects as go
 import plotly.express as ex
 
@@ -22,7 +35,7 @@ from models.cond_encoder import AutoEncoder
 testing = False
 def init_wandb():
     if not testing:
-        run = wandb.init(project='vae', name=config['name'], notes=config['notes'])
+        run = wandb.init(project='ae', name=config['name'], notes=config['notes'])
         artifact = wandb.Artifact('code', type='code')
         artifact.add_file('auto_encoder_curv.py')
         artifact.add_file('losses.py')
@@ -34,7 +47,7 @@ def init_wandb():
         artifact.add_file('convnext/models/convnext_tf.py')
         run.log_artifact(artifact)
 
-uv2mesh = UV2Mesh(res=256)
+uv2mesh = UV2MeshLabel(res=256, is_skirt=config['is_skirt'])
 
 
 g_mean_iou = tf.keras.metrics.MeanIoU(num_classes=7)
@@ -43,6 +56,12 @@ def get_mean_iou(y_true, y_pred):
     g_mean_iou.update_state(y_true, y_pred)
     return g_mean_iou.result()
 
+g_mean_iou_mesh = tf.keras.metrics.MeanIoU(num_classes=7)
+def get_mean_iou_mesh(y_true, y_pred):
+    g_mean_iou_mesh.reset_state()
+    g_mean_iou_mesh.update_state(y_true, y_pred)
+    return g_mean_iou_mesh.result()
+
 # https://github.com/NifTK/NiftyNet/blob/dev/niftynet/layer/loss_segmentation.py
 def dice_dense(prediction, ground_truth, weight_map=None):
     """
@@ -75,7 +94,6 @@ 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)
 
@@ -83,54 +101,62 @@ def train_step(batch):
         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/dice': dice }
+    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):
-    uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
+    pred = model([batch['uv_curvature'], batch['uv_mask']], training=False)
 
-    pred = model([uv_curvature, batch['uv_mask']])
+    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_uv_loss = cce(batch['uv_curvature'], pred)
+    pred_mesh = tf.math.argmax(uv2mesh(pred), axis=-1)
-    ms_ssim = tf.reduce_mean(tf.image.ssim(pred, uv_curvature, max_val=1))
+    mean_iou = get_mean_iou(tf.math.argmax(batch['uv_curvature'], axis=-1), tf.math.argmax(pred, axis=-1))
-    loss =  ce_uv_loss
+    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'])
 
-    mean_iou = get_mean_iou(batch['uv_curvature'], tf.math.argmax(pred, axis=-1))
+    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 }
-    dice = dice_dense(pred, 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 #'ce_mesh_loss': ce_mesh_loss
+
-             }
+def plot_everything(batch: dict, type: str):
+    pred = model([batch['uv_curvature'], batch['uv_mask']], training=False)
+    
+    orig_img = tf.math.argmax(batch['uv_curvature'], axis=-1, output_type=tf.dtypes.int32)
+    pred_img = tf.math.argmax(pred, axis=-1, output_type=tf.dtypes.int32)
+    img = tf.concat((orig_img, pred_img), axis=2).numpy()
+    imgs = { type + '/images': [wandb.Image(img[i,:,:,np.newaxis] / 6.0) for i in range(img.shape[0])] }
+
+    return imgs
 
 
 def train():
+    valid_iterator = iter(dataset_valid)
+
     step = 0
     for batch in dataset_train:
         batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
 
         train_loss = train_step(batch)
 
-        # data_valid = valid_iterator.get_next()
+        data_valid = valid_iterator.get_next()
-        # data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, signature_dynamic))
+        data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, config['signature_dynamic']))
-        valid_loss = {}#valid_step(data_valid)
+        valid_loss = valid_step(data_valid)
-        data_valid = batch
 
         to_log = { **train_loss, **valid_loss}
-        if step % 2500 == 0:
+        if step % 5000 == 0 or (step < 5000 and step % 500 == 0):
-            pred = model([data_valid['uv_curvature'], data_valid['uv_mask']], training=False)
+            imgs_train = plot_everything(batch, 'train')
-            
+            imgs_valid = plot_everything(data_valid, 'valid')
-            orig_img = tf.math.argmax(data_valid['uv_curvature'], axis=-1, output_type=tf.dtypes.int32)
+            to_log = { **to_log, **imgs_train, **imgs_valid }
-            pred_img = tf.math.argmax(pred, axis=-1, output_type=tf.dtypes.int32)
-            img = tf.concat((orig_img, pred_img), axis=2).numpy()
-            # imgs = { 'valid/images': [wandb.Image(img[i,:,:,np.newaxis] / 6.0) for i in range(4)] }
-            imgs = { 'train/images': [wandb.Image(img[i,:,:,np.newaxis] / 6.0) for i in range(img.shape[0])] }
-            to_log = { **to_log, **imgs}
 
         wandb.log(to_log, step=step)  
         print(step)
@@ -174,10 +200,10 @@ if __name__ == '__main__':
 
     factory = DatasetFactory()
 
-    dataset_train = factory.get_dataset(folder=config['train_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']])
+    dataset_train = factory.get_dataset(folder=config['train_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=config['dataset_in_memory'])
-    static_dataset_train = StaticDataset(folder=config['train_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']])
+    static_dataset_train = StaticDataset(folder=config['train_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=True)
-    dataset_valid = factory.get_dataset(folder=config['valid_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'])
+    dataset_valid = factory.get_dataset(folder=config['valid_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=config['dataset_in_memory'])
-    static_dataset_valid = StaticDataset(folder=config['valid_dir'], signature=config['signature_static'])
+    static_dataset_valid = StaticDataset(folder=config['valid_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=True)
 
     model = AutoEncoder(config['cond_ae'])
     optimizer = tf.optimizers.Adam(learning_rate=config['lr'])

configs/config_ae_curv.py

@@ -1,13 +1,12 @@
 from configs.config_base import config
 
-config['gpu'] = '0'
+config['gpu'] = '2'
 config['lr'] = 0.001
-config['name'] = 'ae_curv_9_skirt' #'ae_curv_9'
+config['name'] = 'curv_1_3' #'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['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 | \
+  try ssim loss'
 
 config['cond_ae']['is_label'] = True
-config['cond_ae']['upscale_output'] = False
-
 
 config['signature_dynamic'] = { 'id': (),
                           'uv_curvature': (512, 256, 7),

custom_utils/custom_utils.py

@@ -71,12 +71,12 @@ class UV2Mesh():
         
 
 class UV2MeshLabel():
-    def __init__(self, res=256):
+    def __init__(self, res=256, is_skirt=False):
         super(UV2MeshLabel, self).__init__()
         self.res = res
-        uv_map = np.load('custom_utils/stretched_uv_map_3.npy')
+        uv_map = np.load('custom_utils/stretched_uv_map_' + ('skirt' if is_skirt else '3') + '.npy')
         self.uv_map_discreet = np.floor(uv_map * self.res).astype(np.int32)
-        m2uv = np.load('custom_utils/m2uv.npy', allow_pickle=True).item()
+        m2uv = np.load('custom_utils/m2uv' + ('_skirt' if is_skirt else '') + '.npy', allow_pickle=True).item()
         m2uv = m2uv_array(m2uv)
         self.uv_pixels = self.uv_map_discreet[m2uv]