update curvature CVAE

convnext/models/convnext_tf.py

@@ -369,7 +369,7 @@ class ConvNeXtReverseLabels(tf.keras.Model):
         self.last = Dense(nr_of_classes, activation='softmax')
 
     def call(self, x):
-        for i in reversed(len(depths)):
+        for i in reversed(range(len(self.downsample_layers))):
             x = self.downsample_layers[i](x)
             x = self.stages[i](x)
         x = self.last(x)

cvae_curv.py

@@ -1,5 +1,5 @@
 import os
-os.environ["CUDA_VISIBLE_DEVICES"] = "3"
+os.environ["CUDA_VISIBLE_DEVICES"] = "1"
 import tensorflow as tf
 
 gpus = tf.config.experimental.list_physical_devices('GPU')
@@ -26,10 +26,7 @@ import plotly.graph_objects as go
 import plotly.express as ex
 from PIL import Image
 
-from utils import str_to_bool
+from tf_utils import check_tf_version, allow_memory_growth
-from tf_utils import check_tf_version, allow_memory_growth, split_gpu_for_testing
-from load_models import load_generator, load_discriminator
-from losses import d_logistic, d_logistic_r1_reg, g_logistic_non_saturating, g_logistic_ns_pathreg, edge_loss, uv_mask_loss, bend_loss
 # from losses import collision_penalty as collision_penalty
 from losses import collision_loss as collision_penalty
 from dataset import DatasetFactory, StaticDataset
@@ -38,14 +35,15 @@ from custom_utils.smpl import SMPL
 from models.cvae_curv_model import CVAE
 
 
-name='cvae_curv_2_2_2'
+name='cvae_curv_3'
 def init_wandb():
     run = wandb.init(project='vae', name=name,
             notes='learn offsets. smaller flattened latent vectors. with 25-85k pretraining. smaller batchsize. new implementation for kl loss. \
 no lr schedule. 25 jumpsuits, average reconst_loss before calculation, offsets over average unposed garment. gamma balancing. remove global \
-orientation (in posed body input and in pose_shape parameters.. retrained body encoder with deoriented bodies. lower kl coeff to 0.1. fixed plotting colors\
+orientation (in posed body input and in pose_shape parameters.. r lower kl coeff to 0.1. fixed plotting colors\
-fixed deoriented bodies and need retraining (ae_body_6 with 50k iteration), addent MLP for lowering conditional Z diemnsion\
+addent MLP for lowering conditional Z diemnsion\
-higher starting kl coeff (0.1)')
+higher starting kl coeff (0.1)\
+train on 100 garments')
     artifact = wandb.Artifact('code', type='code')
     artifact.add_file('cvae_curv.py')
     artifact.add_file('models/cvae_curv_model.py')
@@ -161,7 +159,7 @@ def train_step(batch):
     with tf.GradientTape() as tape:
         uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
 
-        pred, mean, logvar, z = model([batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask'], batch['uv_shaped_body']])
+        pred, mean, logvar, z = model([batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask']])
 
         ce_uv_loss = scce(batch['uv_curvature'], pred)
         ms_ssim = tf.reduce_mean(tf.image.ssim(pred, uv_curvature, max_val=1))
@@ -197,7 +195,7 @@ def train_step(batch):
 def valid_step(batch):
     uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
 
-    pred, mean, logvar, z = model([batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask'], batch['uv_shaped_body']])
+    pred, mean, logvar, z = model([batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask']])
 
     ce_uv_loss = scce(batch['uv_curvature'], pred)
     ms_ssim = tf.reduce_mean(tf.image.ssim(pred, uv_curvature, max_val=1))
@@ -220,7 +218,7 @@ def valid_step(batch):
 
 def plot_everything(batch: dict, type: str):
     uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
-    pred, mean, logvar, z = model([batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask'], batch['uv_shaped_body']])
+    pred, mean, logvar, z = model([batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask']])
     
     pred = np.argmax(pred.numpy(), axis=-1)
     mesh_pred = uv2mesh(pred)
@@ -336,7 +334,7 @@ def sample(folder):
         print()
 
 if __name__ == '__main__':
-    # init_wandb()
+    init_wandb()
     allow_memory_growth()
     parser = argparse.ArgumentParser(description='')
     parser.add_argument('--name', type=str)
@@ -363,17 +361,17 @@ if __name__ == '__main__':
                         #  'mesh_static': (14475, 3,),
                          'mesh_mask': (14475,),
                          'uv_mask': (512, 256,),
-                         'uv_shaped_body': (3, 512, 256,),
+                        #  'uv_shaped_body': (3, 512, 256,),
                         }
 
-    dataset_train = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/train', batch_size=batch_size, signature=signature_dynamic)
+    dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100', batch_size=batch_size, signature=signature_dynamic, use_filled=False, resolutions=[512])
-    static_dataset_train = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/train', signature=signature_static)
+    static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', signature=signature_static, use_filled=False, resolutions=[512])
     dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/valid', batch_size=batch_size, signature=signature_dynamic)
     static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/valid', signature=signature_static)
 
 
-    model = CVAE(latent_dim=128, static_model='checkpoints/ae_static_3',
+    model = CVAE(latent_dim=128, static_model='checkpoints/ae_static_4',
-                 body_model='checpoints/ae_body_6',)
+                 body_model='checpoints/ae_body_8',)
 
 
     # lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(0.0001, decay_steps=30000, decay_rate=0.1, staircase=False)
@@ -384,14 +382,13 @@ if __name__ == '__main__':
     valid_iterator = iter(dataset_valid)
     gamma = 1
     balancing_activated_step = 0
-    avg_jumpsuit = np.load('custom_data/jumpsuit_20_avg_uv.npy')
 
 
     # sample('checkpoints/cvae_curv_2')
 
 
     ckpt = tf.train.Checkpoint(model=model)
-    manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + name, max_to_keep=5)
+    manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + name, max_to_keep=3)
 
     step = 0
     for batch in dataset_train:

models/cvae_curv_model.py

@@ -122,21 +122,33 @@ class CVAE(tf.keras.Model):
         return x
 
 
-    def sample(self, uv_template, uv_body_posed, uv_mask, eps=None):
+    def encode_conds(self, uv_template, uv_body_posed):
-        eps = tf.random.normal(shape=(uv_template.shape[0], self.latent_dim))
+        if uv_template.shape[1] == 3:
-
+            uv_template = tf.transpose(uv_template, [0,2,3,1])
-        uv_template = tf.transpose(uv_template, [0,2,3,1])
+        if uv_body_posed.shape[1] == 3:
-        uv_body_posed = tf.transpose(uv_body_posed, [0,2,3,1])
+            uv_body_posed = tf.transpose(uv_body_posed, [0,2,3,1])
 
         z_static = self.static_enc(uv_template)
         z_body = self.body_enc(uv_body_posed)
 
         z_cond = tf.concat([z_static, z_body], axis=-1)
         z_cond = self.cond_encoder(z_cond)
+
+        return z_cond
+
+    def sample(self, uv_template, uv_body_posed, uv_mask=None, noise=None):
+        if noise is None:
+            eps = tf.random.normal(shape=(uv_template.shape[0], self.latent_dim))
+        else:
+            eps = noise
+
+        z_cond = self.encode_conds(uv_template, uv_body_posed)
         z = tf.concat([eps, z_cond], axis=-1)
 
         x = self.decoder(z)
-        x = x * uv_mask[:,:,:,tf.newaxis]
+
+        if uv_mask is not None:
+            x = x * uv_mask[:,:,:,tf.newaxis]
 
         return x
 
@@ -168,17 +180,20 @@ class CVAE(tf.keras.Model):
 
 
     def call(self, inputs, training=None):
-        uv_template, uv_body_posed, uv_curv, uv_mask, uv_shaped_body = inputs
+        uv_template, uv_body_posed, uv_curv, uv_mask = inputs
-        uv_template, uv_body_posed = tf.transpose(uv_template, [0,2,3,1]), tf.transpose(uv_body_posed, [0,2,3,1])
+        if uv_template.shape[1] == 3:
-        z_static = self.static_enc(uv_template)
+            uv_template = tf.transpose(uv_template, [0,2,3,1])
-        z_body = self.body_enc(uv_body_posed)
+        if uv_body_posed.shape[1] == 3:
+            uv_body_posed = tf.transpose(uv_body_posed, [0,2,3,1])
+        if uv_curv.shape[1] == 3:
+            uv_curv = tf.transpose(uv_curv, [0,2,3,1])
+
+        z_cond = self.encode_conds(uv_template, uv_body_posed)
 
         x = tf.concat([uv_template, uv_body_posed, uv_curv], axis=-1)
-
         mean, logvar = self.encoder(x)
         z = self.reparameterize(mean, logvar)
-        z_cond = tf.concat([z_static, z_body], axis=-1)
+
-        z_cond = self.cond_encoder(z_cond)
         z = tf.concat([z, z_cond], axis=-1)
         x = self.decoder(z)