backup

.gitignore

@@ -138,9 +138,10 @@ wandb/
 output/*
 *.zip
 *.out
-models/
+models_old/
 models_archive/
 temp.jpg
 corrupt/
 wandb copy/
 checkpoints/
+temp/

auto_encoder_body.py

@@ -25,16 +25,18 @@ from dataset import DatasetFactory, StaticDataset
 from custom_utils.custom_utils import batch_to_dict
 from custom_utils.smpl import SMPL
 
-name='ae_body_5'
+name='ae_body_7'
 def init_wandb():
-    os.environ["CUDA_VISIBLE_DEVICES"] = "3"
+    os.environ["CUDA_VISIBLE_DEVICES"] = "1"
     run = wandb.init(project='vae', name=name,
             notes='convnext + convnext, add dense with tanh at the end of reverse convnext, larger features. add lr scheduler. fft3d instead of rfft3d. \
     fft2d per channel and separate for top and bottom. fft_loss, all data. added convnext block as bottleneck, same as before, added extra metric. \
     . lowered learning rate because it goes to nan. shallow third stage. slower lr decay. with validation. higher lr rate.\
     two more downscaling than up, so decoder output is resized wto final resolution\
     larger batch_size.\
-    flattened bottleneck with 128 latent_dim\ remove global orientation')
+    flattened bottleneck with 128 latent_dim\ remove global orientation\
+    fixed gender issue with deoriented bodies, again\
+    trained on 100 garments')
     artifact = wandb.Artifact('code', type='code')
     artifact.add_file('auto_encoder_body.py')
     artifact.add_file('losses.py')
@@ -419,8 +421,8 @@ if __name__ == '__main__':
                         #  'uv_shaped_body': (3, 512, 256,),
                         }
 
-    dataset_train = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/train', batch_size=32, signature=signature_dynamic)
-    static_dataset_train = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/train', signature=signature_static)
+    dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100', batch_size=32, signature=signature_dynamic)
+    static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', signature=signature_static)
     dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/valid', batch_size=32, signature=signature_dynamic)
     static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/valid', signature=signature_static)
 

auto_encoder_curv.py

@@ -5,6 +5,9 @@ import wandb
 import random
 import numpy as np
 import tensorflow as tf
+os.environ["CUDA_VISIBLE_DEVICES"] = "2"
+
+
 from custom_utils.custom_utils import UV2Mesh, maskFacesBatch
 from tensorboardX import SummaryWriter
 import plotly.graph_objects as go
@@ -24,14 +27,15 @@ from dataset import DatasetFactory, StaticDataset
 from custom_utils.custom_utils import batch_to_dict
 from custom_utils.smpl import SMPL
 
-name='ae_curvature_6'
+name='ae_curvature_7'
 testing = False
 def init_wandb():
-    os.environ["CUDA_VISIBLE_DEVICES"] = "3"
     if not testing:
-        run = wandb.init(project='vae', name=name, notes='convnext + convnext, 16*8*32, ce loss. flattened latent dim of 128')
+        run = wandb.init(project='vae', name=name, notes='convnext + convnext, 16*8*32, ce loss. flattened latent dim of 128. trained on 100 garments')
         artifact = wandb.Artifact('code', type='code')
         artifact.add_file('auto_encoder_curv.py')
+        artifact.add_file('dataset.py')
+        artifact.add_file('losses.py')
         artifact.add_file('stylegan2_ae/layers/mapping_block.py')
         artifact.add_file('convnext/models/convnext_tf.py')
         run.log_artifact(artifact)
@@ -282,19 +286,21 @@ def train():
 
         train_loss = train_step(batch)
 
-        data_valid = valid_iterator.get_next()
-        data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, signature_dynamic))
-        valid_loss = valid_step(data_valid)
+        # data_valid = valid_iterator.get_next()
+        # data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, signature_dynamic))
+        valid_loss = {}#valid_step(data_valid)
+        data_valid = batch
 
         to_log = { **train_loss, **valid_loss}
-        if step % 1000 == 0:
+        if step % 2500 == 0:
             uv_curvature = tf.one_hot(data_valid['uv_curvature'], 7)
             pred = model([uv_curvature, data_valid['uv_mask']])
             
             orig_img = data_valid['uv_curvature']
             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 = { '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(4)] }
             to_log = { **to_log, **imgs}
 
         wandb.log(to_log, step=step)  
@@ -351,10 +357,10 @@ if __name__ == '__main__':
                         #  'uv_shaped_body': (3, 512, 256,),
                         }
 
-    dataset_train = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/train',
+    dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100',
                                   batch_size=8,
                                   signature=signature_dynamic)
-    static_dataset_train = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/train',
+    static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100',
                                    signature=signature_static)
     dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/valid',
                                   batch_size=8,

auto_encoder_static.py

@@ -24,9 +24,9 @@ from dataset import DatasetFactory, StaticDataset
 from custom_utils.custom_utils import batch_to_dict
 from custom_utils.smpl import SMPL
 
-name='ae_static_3'
+name='ae_static_4'
 def init_wandb():
-    os.environ["CUDA_VISIBLE_DEVICES"] = "2"
+    os.environ["CUDA_VISIBLE_DEVICES"] = "0"
     run = wandb.init(project='vae', name=name,
             notes='convnext + convnext, add dense with tanh at the end of reverse convnext, larger features. add lr scheduler. fft3d instead of rfft3d. \
     fft2d per channel and separate for top and bottom. fft_loss, all data. added convnext block as bottleneck, same as before, added extra metric. \
@@ -389,13 +389,13 @@ def train():
         batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic))
         train_loss = train_step(batch)
 
-        data_valid = valid_iterator.get_next()
-        data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, signature_dynamic))
-        valid_loss = valid_step(data_valid)
+        # data_valid = valid_iterator.get_next()
+        # data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, signature_dynamic))
+        valid_loss = {}#valid_step(data_valid)
 
         if step % 1000 == 0:
             train_figs = plot_everything(batch, 'train')
-            valid_figs = plot_everything(data_valid, 'valid')
+            valid_figs = {}#plot_everything(data_valid, 'valid')
             wandb.log({ **train_loss, **valid_loss, **train_figs, **valid_figs }, step=step)
         else:
             wandb.log({ **train_loss, **valid_loss }, step=step)
@@ -432,8 +432,8 @@ if __name__ == '__main__':
                         #  'uv_shaped_body': (3, 512, 256,),
                         }
 
-    dataset_train = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/train', batch_size=8, signature=signature_dynamic)
-    static_dataset_train = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/train', signature=signature_static)
+    dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100', batch_size=8, signature=signature_dynamic)
+    static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', signature=signature_static)
     dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/valid', batch_size=8, signature=signature_dynamic)
     static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/valid', signature=signature_static)
 

custom_utils/custom_utils.py

@@ -197,3 +197,9 @@ def sparse_to_tensor(L):
     data = L.data.astype(dtype=np.float32)
     shape = L.shape
     return tf.SparseTensor(idx, data, shape)
+
+
+def load_model(model, checkpoint):
+    ckpt = tf.train.Checkpoint(model=model)
+    manager = tf.train.CheckpointManager(ckpt, checkpoint, max_to_keep=1)
+    ckpt.restore(manager.latest_checkpoint)

custom_utils/smpl.py

@@ -69,6 +69,7 @@ class SMPL:
         self.nn = np.array(sio.loadmat('custom_utils/NN.mat'))
         self.to_keep = sio.loadmat(os.path.join('custom_utils/keep_remove.mat'))['keep'].reshape(-1)
 
+        self.weights_all = self.weights
         self.weights = tf.gather(self.weights, self.to_keep, axis=1)
 
     def Q2R(self, w, x, y, z):
@@ -132,7 +133,7 @@ class SMPL:
         v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[1], axes=[[1], [2]]) + tf.reshape(self.v_template[1], shape=(1,27578,3))) # male
         v_shaped = tf.stack(v_shaped, axis=-1)
 
-        def single_run(x, J_regressor=self.J_regressor, posedirs=self.posedirs, weights=self.weights, parent=self.parent):
+        def single_run(x, J_regressor=self.J_regressor, posedirs=self.posedirs, weights=self.weights_all, parent=self.parent):
             R = x[0]
             v_shaped = x[1]
             g = x[2]
@@ -160,7 +161,7 @@ class SMPL:
 
             return v, tf.squeeze(J), v_posed
 
-        gender = tf.cast(tf.greater_equal(betas[:, 10] , 0), tf.int32)
+        gender = tf.cast(tf.greater_equal(betas[:, 10] , 0.5), tf.int32)
         T_out, J_out, v_posed = tf.map_fn(single_run, (R, v_shaped, gender), (tf.float32, tf.float32, tf.float32))
         
         J0 = tf.reshape(J_out[:,0,:], (-1, 1, 3))
@@ -215,7 +216,7 @@ class SMPL:
 
             return v, tf.squeeze(J)
 
-        gender = tf.cast(tf.greater_equal(betas[:, 10] , 0), tf.int32)
+        gender = tf.cast(tf.greater_equal(betas[:, 10] , 0.5), tf.int32)
         T_out, J_out = tf.map_fn(single_run, (R, v_shaped, gender, v_cloth), (tf.float32, tf.float32))
         
         J0 = tf.reshape(J_out[:,0,:], (-1, 1, 3))
@@ -235,7 +236,7 @@ class SMPL:
             J = tf.matmul(J_regressor[g], v_shaped[:, :, g])
             return J
 
-        gender = tf.cast(tf.greater(betas[:, 10] , 0), tf.int32)
+        gender = tf.cast(tf.greater(betas[:, 10] , 0.5), tf.int32)
         J = tf.map_fn(single_run, (v_shaped, gender), (tf.float32))
         J0 = tf.reshape(J[:,0,:], (-1, 1, 3))
         
@@ -274,7 +275,7 @@ class SMPL:
 
             return v, tf.squeeze(J)
 
-        gender = tf.cast(tf.greater(betas[:, 10] , 0), tf.int32)
+        gender = tf.cast(tf.greater(betas[:, 10] , 0.5), tf.int32)
         T_out, J_out = tf.map_fn(single_run, (R, J, gender, v_cloth), (tf.float32, tf.float32))
         
         J0 = tf.reshape(J_out[:,0,:], (-1, 1, 3))

cvae.py

@@ -1,5 +1,5 @@
 import os
-os.environ["CUDA_VISIBLE_DEVICES"] = "2"
+os.environ["CUDA_VISIBLE_DEVICES"] = "3"
 import tensorflow as tf
 
 gpus = tf.config.experimental.list_physical_devices('GPU')
@@ -24,6 +24,7 @@ from custom_utils.custom_utils import UV2Mesh, maskFacesBatch, _maskFaces
 from tensorboardX import SummaryWriter
 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, split_gpu_for_testing
@@ -31,26 +32,22 @@ 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 convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block
-from stylegan2_ae.generator import Generator
-from custom_utils.graph_utils import Reconstructor_V4_Decoder as GNNDecoder, sparse_to_tensor
-from custom_utils.graph_utils import Reconstructor_v2 as SimpleReconstructor
+from custom_utils.custom_utils import batch_to_dict, load_model
 from dataset import DatasetFactory, StaticDataset
-from custom_utils.custom_utils import batch_to_dict
 from custom_utils.smpl import SMPL
-from auto_encoder_body import AutoEncoder as BodyAE
-from auto_encoder_curv import AutoEncoder as CurvatureAE
-from auto_encoder_static import AutoEncoder as StaticAE
+from models.cvae import CVAE
 
 
-name='cvae_10_2'
+name='cvae_15'
 def init_wandb():
     run = wandb.init(project='vae', name=name,
             notes='learn offsets. smaller flattened latent vectors. with 25-85k pretraining. smaller batchsize. fixed balancing of reconst and kl loss for baseline. \
-new implementation for kl loss. no lr schedule. 25 jumpsuits, average reconst_loss before calculation, linearly increasing balancing, up to 1000 from 250.\
-offsets over average unposed garment. gamma balancing. remove global orientation (in posed body input and in pose_shape parameters. larger batchsize\
-add edge and collision (fixed) loss. rebalanced l1 loss. increasing edge and coll coeff. retrained body encoder with deoriented bodies \
-don\'t freeze encoders, lower lr')
+new implementation for kl loss. no lr schedule. 25 jumpsuits, average reconst_loss before calculation, linearly increasing balancing, up to 1 from 0.1. \
+offsets over average unposed garment. remove global orientation (in posed body input and in pose_shape parameters. \
+retrained body encoder with deoriented bodies \
+use mean l2 loss. small batch. l2 * 1000, balancing to 1 from 0.01\
+fixed gender issue and retrained body encoder\
+remove shape from unposed garment')
     artifact = wandb.Artifact('code', type='code')
     artifact.add_file('cvae.py')
     artifact.add_file('dataset.py')
@@ -61,14 +58,14 @@ don\'t freeze encoders, lower lr')
 uv2mesh = UV2Mesh(res=256)
 
 def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, collision_color):
-    mesh_pred_unposed[:,1] += 1
-    collision[:,1] += 2
+    mesh_pred_unposed[:,2] += 1
+    collision[:,2] += 2
     fig = go.Figure(data=[
         go.Mesh3d(
             # 8 vertices of a cube
             x=mesh_gt[:,0],
-            y=mesh_gt[:,1],
-            z=mesh_gt[:,2],
+            y=mesh_gt[:,2],
+            z=mesh_gt[:,1],
             # i, j and k give the vertices of triangles
             i = faces[:,0],
             j = faces[:,1],
@@ -77,8 +74,8 @@ def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, colli
         go.Mesh3d(
             # 8 vertices of a cube
             x=mesh_pred_unposed[:,0],
-            y=mesh_pred_unposed[:,1],
-            z=mesh_pred_unposed[:,2],
+            y=mesh_pred_unposed[:,2],
+            z=mesh_pred_unposed[:,1],
             # i, j and k give the vertices of triangles
             i = faces[:,0],
             j = faces[:,1],
@@ -87,8 +84,8 @@ def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, colli
         go.Mesh3d(
             # 8 vertices of a cube
             x=collision[:,0],
-            y=collision[:,1],
-            z=collision[:,2],
+            y=collision[:,2],
+            z=collision[:,1],
             # i, j and k give the vertices of triangles
             i = collision_faces[:,0],
             j = collision_faces[:,1],
@@ -108,151 +105,7 @@ def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, colli
 
     return fig
 
-class Bottleneck(tf.keras.Model):
 
-    def __init__(self, dim):
-        super(Bottleneck, self).__init__()
-
-        self.bottleneck = tf.keras.Sequential(
-                [Block(dim=dim, drop_path=0,
-                       layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
-            )
-
-    def call(self, inputs, training=None):
-        x = inputs
-
-        x = self.bottleneck(x, training=training)
-        
-        return x
-
-
-class Encoder(tf.keras.Model):
-    
-    def __init__(self, latent_dim):
-        super(Encoder, self).__init__()
-
-        self.encoder = ConvNeXt(depths=[3, 3, 9, 3], dims=[32, 64, 128, 256], include_top=False)
-        self.flatten = tf.keras.layers.Flatten()
-        self.mean = tf.keras.layers.Dense(latent_dim)
-        self.log_var = tf.keras.layers.Dense(latent_dim)
-
-    def call(self, inputs, training=None):
-        x = inputs
-
-        x = self.encoder(x)
-        x = self.flatten(x)
-        mean = self.mean(x)
-        log_var = self.log_var(x)
-
-        return mean, log_var
-
-
-class Decoder(tf.keras.Model):
-    
-    def __init__(self):
-        super(Decoder, self).__init__()
-
-        self.dense = tf.keras.layers.Dense(16 * 8 * 80)
-        self.reshape = tf.keras.layers.Reshape((16, 8, 80))
-        self.decoder = ConvNeXtReverse(depths=list(reversed([3, 3, 9, 3])), dims=list(reversed([256, 128, 64, 16])))
-
-
-    def call(self, inputs, training=None):
-        x = inputs
-
-        x = self.dense(x)
-        x = self.reshape(x)
-        x = self.decoder(x)
-
-        return x
-
-
-class CVAE(tf.keras.Model):
-    
-    def __init__(self, latent_dim=128):
-        super(CVAE, self).__init__()
-        self.latent_dim = latent_dim
-
-        model : StaticAE = StaticAE()
-        load_model(model, 'checkpoints/ae_static_3')
-        self.static_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
-        # self.static_enc.trainable = False
-
-        model : BodyAE = BodyAE()
-        load_model(model, 'checpoints/ae_body_5')
-        self.body_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
-        # self.body_enc.trainable = False
-
-        model : CurvatureAE = CurvatureAE()
-        load_model(model, 'checpoints/ae_curvature_6')
-        self.curv_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
-        # self.curv_enc.trainable = False
-
-        self.encoder = Encoder(latent_dim)
-        self.decoder = Decoder()
-
-
-    # @tf.function
-    def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask, eps=None):
-        if eps is None:
-            eps = tf.random.normal(shape=(4, self.latent_dim))
-
-        uv_template = tf.tile(tf.transpose(uv_template, [1,2,0])[tf.newaxis,:,:,:], [4,1,1,1])
-        uv_body_posed = tf.tile(tf.transpose(uv_body_posed, [1,2,0])[tf.newaxis,:,:,:], [4,1,1,1])
-        uv_curv = tf.tile(uv_curv[tf.newaxis,:,:,:], [4,1,1,1])
-        uv_mask = tf.tile(uv_mask[tf.newaxis,:,:], [4,1,1])
-
-        # uv_template, uv_body_posed = tf.transpose(uv_template, [0,2,3,1]), 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_curv = self.curv_enc(uv_curv)
-
-        z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
-        z = tf.concat([eps, z_cond], axis=-1)
-
-        x = self.decoder(z)
-        x = x + avg_jumpsuit[tf.newaxis,:,:,:]
-        x = x * uv_mask[:,:,:,tf.newaxis]
-
-        return x
-
-
-    def reparameterize(self, mean, logvar):
-        eps = tf.random.normal(shape=mean.shape)
-        return eps * tf.exp(logvar * .5) + mean
-
-
-    def decode(self, z, apply_sigmoid=False):
-        logits = self.decoder(z)
-        if apply_sigmoid:
-            probs = tf.sigmoid(logits)
-            return probs
-        return logits
-
-
-    def call(self, inputs, training=None):
-        uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask, uv_shaped_body = inputs
-        uv_unposed, uv_template, uv_body_posed = tf.transpose(uv_unposed, [0,2,3,1]), tf.transpose(uv_template, [0,2,3,1]), 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_curv = self.curv_enc(uv_curv)
-
-        x = tf.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], axis=-1)
-
-        mean, logvar = self.encoder(x)
-        z = self.reparameterize(mean, logvar)
-        z = tf.concat([z, z_static, z_body, z_curv], axis=-1)
-        x = self.decoder(z)
-
-
-        # x = tf.transpose(x, [0,3,1,2])
-        # x = x * uv_mask[:,tf.newaxis,:,:]
-        # x = tf.image.resize(x, (512, 256), tf.image.ResizeMethod.LANCZOS5)
-        # x = x + tf.transpose(uv_shaped_body, [0,2,3,1])
-        x = x + avg_jumpsuit[tf.newaxis,:,:,:]
-        x = x * uv_mask[:,:,:,tf.newaxis]
-
-        return x, mean, logvar, z
 
 
 def log_normal_pdf(sample, mean, logvar, raxis=1):
@@ -268,7 +121,7 @@ def kl_loss(mean, logvar):
     return kl_loss
 
 
-@tf.function
+# @tf.function
 def train_step(batch):
 
     with tf.GradientTape() as tape:
@@ -280,10 +133,15 @@ def train_step(batch):
         mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred)
         mesh_orig = batch['mesh_unposed']
 
+        shape = batch['pose_shape'][:, 72:]
+        gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
+        shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
+        mesh_pred_posed = mesh_pred_posed + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
 
         e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
         c_loss = tf.reduce_mean(collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']))
         l1_loss = tf.reduce_sum(tf.reduce_sum((tf.math.abs(tf.transpose(batch['uv_unposed'], [0,2,3,1]) - pred)), axis=-1) * batch['uv_mask'], axis=[1, 2])
+        l2_loss = 1e3 * tf.reduce_mean(tf.reduce_sum(((tf.transpose(batch['uv_unposed'], [0,2,3,1]) - pred) ** 2), axis=-1) * batch['uv_mask'], axis=[1, 2])
         # ms_ssim = tf.image.ssim_multiscale(tf.transpose(pred, [0,2,3,1]) + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=2)
         ms_ssim = tf.image.ssim_multiscale(pred + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=3)
         # ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale((tf.transpose(pred, [0,2,3,1]) + 1) * batch['uv_mask'][:,:,:,tf.newaxis], (tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1) * batch['uv_mask'][:,:,:,tf.newaxis], max_val=2))
@@ -292,17 +150,18 @@ def train_step(batch):
         fft_loss_pc, fft_loss_pc_low, fft_loss_pc_high = 0,0,0 #fft_loss_per_channel(pred, batch['uv_unposed'])
         tv_loss = tf.image.total_variation(pred)
 
-        reconstr_loss = l1_loss / 4
-        reconstr_loss = tf.reduce_mean(reconstr_loss) + reg_coeff *(e_loss + c_loss)
+        reconstr_loss = l2_loss
+        reconstr_loss = tf.reduce_mean(reconstr_loss)
         # logpx_z = -reconstr_loss
         # logpz = log_normal_pdf(z[:,:model.latent_dim], 0., 0.)
         # logqz_x = log_normal_pdf(z[:,:model.latent_dim], mean, logvar)
         # loss = -tf.reduce_mean(logpx_z + balancing_coeff * (logpz - logqz_x))
         kl = kl_loss(mean, logvar)
-        loss = reconstr_loss / (2 * (gamma ** 2)) + balancing_coeff * kl
+        loss = reconstr_loss + balancing_coeff * kl
         # loss = reconstr_loss + balancing_coeff * kl
         ms_ssim = tf.reduce_mean(ms_ssim)
         l1_loss = tf.reduce_mean(l1_loss)
+        l2_loss = tf.reduce_mean(l2_loss)
         fft_loss = tf.reduce_mean(fft_loss)
         tv_loss = tf.reduce_mean(tv_loss)
         mean = tf.reduce_mean(mean)
@@ -320,7 +179,7 @@ def train_step(batch):
              'train/fft_loss': fft_loss, 'train/euclid_error_reconst': metric_2, 'train/rfft_loss': rfft_loss, 'train/fft_oc_low': fft_loss_pc_low, 'train/fft_pc_high': fft_loss_pc_high,
              'train/fft_loss_per_channel': fft_loss_pc, 'train/euclid_error_posed': euclid_posed, 'train/reconstr_loss': reconstr_loss, 'train/tv_loss': tv_loss,
              'train/mean': mean, 'train/logvar': logvar, 'train/kl_loss': kl, 'train/reconst_loss_norm': reconstr_loss / (2 * (gamma ** 2)), 'train/var': np.e ** logvar,
-             'train/reg_coeff': reg_coeff }
+             'train/reg_coeff': reg_coeff, 'train/l2_loss': l2_loss }
         
 
 
@@ -384,7 +243,12 @@ def plot_everything(batch: dict, type: str):
     
     body = batch['mesh_body_posed'].numpy()
     mesh_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred).numpy()
-    garment_posed = [mesh_posed[i, mask[i], :] for i in range(body.shape[0])]
+
+    shape = batch['pose_shape'][:, 72:]
+    gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
+    shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
+
+    garment_posed = [mesh_posed[i, mask[i], :] + shape_offsets.numpy()[i, mask[i], :] for i in range(body.shape[0])]
     collision = [np.concatenate((body[i], garment_posed[i]), axis=0) for i in range(len(garment_posed))]
     collision_faces = [np.concatenate((F, faces[i] + 14475), axis=0) for i in range(len(garment_posed))]
     collision_color = [np.concatenate((np.ones_like(body[i], dtype=np.int32)*100, np.ones_like(garment_posed[i], dtype=np.int32)*230), axis=0) for i in range(len(garment_posed))]
@@ -397,25 +261,25 @@ def plot_everything(batch: dict, type: str):
     return { **figs, **{ type + '/images': imgs, type + '/mean_hist': mean_hist, type + '/logvar_hist': logvar_hist } }
 
 
-def load_model(model, checkpoint):
-    ckpt = tf.train.Checkpoint(model=model)
-    manager = tf.train.CheckpointManager(ckpt, checkpoint, max_to_keep=1)
-    ckpt.restore(manager.latest_checkpoint)
 
 
 def sample(folder):
+    # run = wandb.init(project='vae_plots', name=folder.split('/')[-1])
     load_model(model,folder)
+    count = 0
     for batch in dataset_train:
         batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic))
         uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
+        nr_of_samples=2
 
-        pred = model.sample(batch['uv_static'][0], batch['uv_body_posed'][0], uv_curvature[0], batch['uv_mask'][0])
+        pred = model.sample(batch['uv_static'][0], batch['uv_body_posed'][0], uv_curvature[0], batch['uv_mask'][0], nr_of_samples=nr_of_samples)
         pose, shape = batch['pose_shape'][0, :72], batch['pose_shape'][0, 72:]
-        pose, shape = tf.tile(pose[tf.newaxis,:], [4,1]), tf.tile(shape[tf.newaxis,:], [4,1])
+        pose, shape = tf.tile(pose[tf.newaxis,:], [nr_of_samples,1]), tf.tile(shape[tf.newaxis,:], [nr_of_samples,1])
 
         mesh_pred = uv2mesh(tf.transpose(pred, [0,3,1,2]))
         mesh_posed = smpl.run_garment(pose, shape, mesh_pred).numpy()
         mesh_orig = batch['mesh_posed'][0].numpy()
+        mesh_orig_reposed = smpl.run_garment(pose[:1], shape[:1], batch['mesh_unposed'][:1]).numpy()[0]
 
         body = batch['mesh_body_posed'][0].numpy()
         body_color = np.ones_like(body, dtype=np.int32)*100
@@ -424,72 +288,110 @@ def sample(folder):
         mask = batch['mesh_mask'].numpy().astype(bool)[0]
         _, faces = _maskFaces(np.ones((14475, 3)), F, mask)
 
-        batch = 0
+
+        pred_zero = model.sample(batch['uv_static'][0], batch['uv_body_posed'][0], uv_curvature[0], batch['uv_mask'][0], nr_of_samples=nr_of_samples, zero_latent=True)
+        mesh_pred_zero = uv2mesh(tf.transpose(pred_zero, [0,3,1,2]))
+        mesh_posed_zero = smpl.run_garment(pose, shape, mesh_pred_zero).numpy()
+
+
+
 
         meshes = []
         meshes.append(go.Mesh3d(
                     # 8 vertices of a cube
-                    x=mesh_orig[mask,0],
-                    y=mesh_orig[mask,1],
-                    z=mesh_orig[mask,2],
+                    x=mesh_orig_reposed[mask,0],
+                    y=mesh_orig_reposed[mask,2],
+                    z=mesh_orig_reposed[mask,1],
                     # i, j and k give the vertices of triangles
                     i = faces[:,0],
                     j = faces[:,1],
                     k = faces[:,2]
                     ))
 
-        for i in range(4):
+        for i in range(nr_of_samples):
             x = mesh_posed[i][mask, :].copy()
-            x[:, 1] += 1 * (i + 1)
+            # x[:, 2] += 1 * (i + 1)
+            x[:, 2] += 1
             meshes.append(go.Mesh3d(
                         # 8 vertices of a cube
                         x=x[:,0],
-                        y=x[:,1],
-                        z=x[:,2],
+                        y=x[:,2],
+                        z=x[:,1],
                         # i, j and k give the vertices of triangles
                         i = faces[:,0],
                         j = faces[:,1],
                         k = faces[:,2]
                         ))
 
-        for i in range(5):
-            x = body.copy()
-            x[:, 1] += 1 * i
+        
+        for i in range(nr_of_samples):
+            x = mesh_posed_zero[i][mask, :].copy()
+            # x[:, 2] += 1 * (i + 1)
+            x[:, 2] += 2
             meshes.append(go.Mesh3d(
                         # 8 vertices of a cube
                         x=x[:,0],
-                        y=x[:,1],
-                        z=x[:,2],
+                        y=x[:,2],
+                        z=x[:,1],
                         # i, j and k give the vertices of triangles
-                        i = F[:,0],
-                        j = F[:,1],
-                        k = F[:,2],
-                        vertexcolor=body_color
+                        i = faces[:,0],
+                        j = faces[:,1],
+                        k = faces[:,2]
                         ))
 
+        # for i in range(5):
+        #     x = body.copy()
+        #     x[:, 2] += 1 * i
+        #     meshes.append(go.Mesh3d(
+        #                 # 8 vertices of a cube
+        #                 x=x[:,0],
+        #                 y=x[:,2],
+        #                 z=x[:,1],
+        #                 # i, j and k give the vertices of triangles
+        #                 i = F[:,0],
+        #                 j = F[:,1],
+        #                 k = F[:,2],
+        #                 vertexcolor=body_color
+        #                 ))
+
 
         fig = go.Figure(data=meshes)
 
-        fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=3, z=1))
+        fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=2, z=1))
+        # fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=nr_of_samples-1, z=1))
         fig.update_layout(
             scene = dict(
                 xaxis = dict(nticks=4, range=[-1,1],),
-                            yaxis = dict(nticks=4, range=[-1,5],),
+                            # yaxis = dict(nticks=4, range=[-1,nr_of_samples+1],),
+                            yaxis = dict(nticks=4, range=[-1,3],),
                             zaxis = dict(nticks=4, range=[-1,1],),)
                             )
 
+        # wandb.log({'plot_' + str(count): fig})
         fig.show()
 
+        pred = pred.numpy()
+        frames = [Image.fromarray(((pred[i] + 1) * 128).astype(np.int8), mode='RGB') for i in range(pred.shape[0])]
+        frames[0].save("temp/garment_1.gif", format="GIF", append_images=frames[1:],
+               save_all=True, duration=500, loop=0)
+
+
+        if count > 10:
+            break
+
+        print(count)
+        count += 1
+
         print()
 
 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())
 
-    batch_size = 12
+    batch_size = 4
     
     smpl = SMPL()
 
@@ -511,27 +413,30 @@ if __name__ == '__main__':
                          '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)
-    static_dataset_train = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/train', signature=signature_static)
+    dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100', batch_size=batch_size, signature=signature_dynamic)
+    static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', signature=signature_static)
     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()
+    model = CVAE(avg_garment_path='custom_data/jumpsuit_20_avg_uv.npy',
+                 latent_dim=128, static_model='checkpoints/ae_static_3',
+                 body_model='checpoints/ae_body_6',
+                 curvature_model='checpoints/ae_curvature_6')
 
 
     # lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(0.0001, decay_steps=30000, decay_rate=0.1, staircase=False)
     optimizer = tf.optimizers.Adam(learning_rate=0.0001)
-    balancing_coeff = 250
+    balancing_coeff = 0.0001
     reg_coeff = 1
     folder = 'models/ae_logs/'
     valid_iterator = iter(dataset_valid)
     gamma = 1
     balancing_activated_step = 0
-    avg_jumpsuit = np.load('custom_data/jumpsuit_20_avg_uv.npy')
+    # avg_jumpsuit = np.load('custom_data/jumpsuit_20_avg_uv.npy')
 
 
-    sample('checkpoints/cvae_10')
+    # sample('checkpoints/cvae_13')
 
 
     ckpt = tf.train.Checkpoint(model=model)
@@ -561,9 +466,9 @@ if __name__ == '__main__':
         if train_loss['train/euclid_error'] < 0.01:
             if balancing_activated_step == 0:
                 balancing_activated_step = step
-            balancing_coeff = balancing_coeff + 1000 / 30000
-        if train_loss['train/euclid_error'] < 0.015:
-            reg_coeff = reg_coeff + 100 / 15000
+            balancing_coeff = balancing_coeff + 1 / 15000
+        # if train_loss['train/euclid_error'] < 0.015:
+        #     reg_coeff = reg_coeff + 100 / 15000
         # if step > 15000 and balancing_coeff <= 1e-3:
         #     balancing_coeff = 1e-2
         # if step > 20000 and balancing_coeff <= 1e-3:

cvae_curv.py

@@ -0,0 +1,440 @@
+import os
+os.environ["CUDA_VISIBLE_DEVICES"] = "3"
+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 cmapy as cm
+import numpy as np
+import scipy as sp
+from custom_utils.custom_utils import UV2Mesh, UV2MeshLabel, maskFacesBatch, _maskFaces
+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, 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
+from custom_utils.custom_utils import batch_to_dict
+from custom_utils.smpl import SMPL
+from models.cvae_curv import CVAE
+
+
+name='cvae_curv_2_2_2'
+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\
+fixed deoriented bodies and need retraining (ae_body_6 with 50k iteration), addent MLP for lowering conditional Z diemnsion\
+higher starting kl coeff (0.1)')
+    artifact = wandb.Artifact('code', type='code')
+    artifact.add_file('cvae_curv.py')
+    artifact.add_file('dataset.py')
+    artifact.add_file('losses.py')
+    artifact.add_file('convnext/models/convnext_tf.py')
+    run.log_artifact(artifact)
+
+uv2mesh = UV2MeshLabel(res=256)
+
+def plot_3d(curv_gt, curv_pred, faces, mesh):
+    mesh_2 = mesh.copy()
+    mesh_2[:,1] += 1
+
+    colors = np.array([cm.color('jet', 1 - i) for i in np.linspace(0, 1, 6+1)])
+    v_color_1 = np.zeros((mesh.shape[0], 3))
+    v_color_2 = np.zeros((mesh.shape[0], 3))
+    for i in range(6+1):
+        v_color_1[curv_gt == i, :] = colors[i]
+        v_color_2[curv_pred == i, :] = colors[i]
+
+    fig = go.Figure(data=[
+        go.Mesh3d(
+            # 8 vertices of a cube
+            x=mesh[:,0],
+            y=mesh[:,1],
+            z=mesh[:,2],
+            # i, j and k give the vertices of triangles
+            i = faces[:,0],
+            j = faces[:,1],
+            k = faces[:,2],
+            vertexcolor=v_color_1,
+            text='gt'
+        ),
+        go.Mesh3d(
+            # 8 vertices of a cube
+            x=mesh_2[:,0],
+            y=mesh_2[:,1],
+            z=mesh_2[:,2],
+            # i, j and k give the vertices of triangles
+            i = faces[:,0],
+            j = faces[:,1],
+            k = faces[:,2],
+            vertexcolor=v_color_2,
+            text='pred'
+        )
+    ])
+
+    fig.update_layout(scene_aspectmode='manual',
+                    scene_aspectratio=dict(x=1, y=2, z=1))
+    fig.update_layout(
+        scene = dict(
+            xaxis = dict(nticks=4, range=[-1,1],),
+                        yaxis = dict(nticks=4, range=[-1,3],),
+                        zaxis = dict(nticks=4, range=[-1,1],),)
+                        )
+
+    return fig
+
+
+def log_normal_pdf(sample, mean, logvar, raxis=1):
+    log2pi = tf.math.log(2. * np.pi)
+    res = tf.reduce_sum(-.5 * ((sample - mean) ** 2. * tf.exp(-logvar) + logvar + log2pi), axis=raxis)
+    return res
+
+
+def kl_loss(mean, logvar):
+    kl_loss = 1 + logvar - tf.square(mean) - tf.exp(logvar)
+    kl_loss = tf.reduce_mean(kl_loss)
+    kl_loss *= -0.5
+    return kl_loss
+
+
+
+g_mean_iou = tf.keras.metrics.MeanIoU(num_classes=7)
+def get_mean_iou(y_true, y_pred):
+    g_mean_iou.reset_state()
+    g_mean_iou.update_state(y_true, y_pred)
+    return g_mean_iou.result()
+
+# https://github.com/NifTK/NiftyNet/blob/dev/niftynet/layer/loss_segmentation.py
+def dice_dense(prediction, ground_truth, weight_map=None):
+    """
+    Computing mean-class Dice similarity.
+    :param prediction: last dimension should have ``num_classes``
+    :param ground_truth: segmentation ground truth (encoded as a binary matrix)
+        last dimension should be ``num_classes``
+    :param weight_map:
+    :return: ``1.0 - mean(Dice similarity per class)``
+    """
+
+    if weight_map is not None:
+        raise NotImplementedError
+    prediction = tf.cast(prediction, dtype=tf.float32)
+    ground_truth = tf.cast(ground_truth, dtype=tf.float32)
+    ground_truth = tf.reshape(ground_truth, prediction.shape)
+    # computing Dice over the spatial dimensions
+    reduce_axes = list(range(len(prediction.shape) - 1))
+    dice_numerator = 2.0 * tf.reduce_sum(
+        prediction * ground_truth, axis=reduce_axes)
+    dice_denominator = \
+        tf.reduce_sum(tf.square(prediction), axis=reduce_axes) + \
+        tf.reduce_sum(tf.square(ground_truth), axis=reduce_axes)
+
+    epsilon = 0.00001
+
+    dice_score = (dice_numerator + epsilon) / (dice_denominator + epsilon)
+    return tf.reduce_mean(dice_score)
+
+
+@tf.function
+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']])
+
+        ce_uv_loss = scce(batch['uv_curvature'], pred)
+        ms_ssim = tf.reduce_mean(tf.image.ssim(pred, uv_curvature, max_val=1))
+
+        reconstr_loss = ce_uv_loss
+        reconstr_loss = tf.reduce_mean(reconstr_loss)
+        kl = kl_loss(mean, logvar)
+        loss = reconstr_loss + balancing_coeff * kl
+        ms_ssim = tf.reduce_mean(ms_ssim)
+        mean = tf.reduce_mean(mean)
+        logvar = tf.reduce_mean(logvar)
+
+
+    # mesh_pred = uv2mesh(tf.tile(tf.math.argmax(pred, axis=-1, output_type=tf.dtypes.int32)[:,:,:,tf.newaxis], [1,1,1,3]))[:,:,0]
+    # mesh_orig = tf.one_hot(batch['mesh_curvature'] + tf.cast(batch['mesh_mask'], tf.int32), 7)
+    # ce_mesh_loss = scce(mesh_orig, mesh_pred)
+
+    mean_iou = get_mean_iou(batch['uv_curvature'], tf.math.argmax(pred, axis=-1))
+    dice = dice_dense(pred, 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,  #'ce_mesh_loss': ce_mesh_loss
+             'train/reconstr_loss': reconstr_loss,
+             'train/mean': mean, 'train/logvar': logvar, 'train/kl_loss': kl, 'train/reconst_loss_norm': reconstr_loss / (2 * (gamma ** 2)), 'train/var': np.e ** logvar,
+             'train/reg_coeff': reg_coeff}
+
+
+        
+
+
+@tf.function
+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']])
+
+    ce_uv_loss = scce(batch['uv_curvature'], pred)
+    ms_ssim = tf.reduce_mean(tf.image.ssim(pred, uv_curvature, max_val=1))
+    loss =  ce_uv_loss
+
+    mean_iou = get_mean_iou(batch['uv_curvature'], tf.math.argmax(pred, axis=-1))
+    dice = dice_dense(pred, uv_curvature)
+
+    reconstr_loss = tf.reduce_mean(ce_uv_loss)
+    kl = kl_loss(mean, logvar)
+    loss = reconstr_loss + balancing_coeff * kl
+    ms_ssim = tf.reduce_mean(ms_ssim)
+    mean = tf.reduce_mean(mean)
+    logvar = tf.reduce_mean(logvar)
+
+    return { 'valid/loss': loss, 'valid/ce_uv_loss': ce_uv_loss, 'valid/ssim': ms_ssim, 'valid/mean_iou': mean_iou, 'valid/dice': dice,
+            'valid/ms_ssim': ms_ssim, 'valid/reconstr_loss': reconstr_loss,
+             'valid/mean': mean, 'valid/logvar': logvar, 'valid/kl_loss': kl, 'valid/var': np.e ** logvar }
+
+
+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 = np.argmax(pred.numpy(), axis=-1)
+    mesh_pred = uv2mesh(pred)
+    mesh_orig = batch['mesh_curvature']
+    mesh_base = batch['mesh_posed']
+    
+    orig_img = batch['uv_curvature']
+    pred_img = pred
+    img = tf.concat((orig_img, pred_img), axis=2)
+    imgs = [wandb.Image(img[i].numpy().astype(np.float32) / 6.0) for i in range(4)]
+
+    F = np.load("custom_utils/smpl_faces.npy")
+    mask = batch['mesh_mask'].numpy().astype(bool)
+    _, faces = maskFacesBatch(np.zeros((4, 14475, 3)), F, mask)
+    gt_curv = [mesh_orig[i].numpy()[mask[i]] + 1 for i in range(pred.shape[0])]
+    pred_curv = [mesh_pred[i].numpy()[mask[i]] for i in range(pred.shape[0])]
+    mesh = [mesh_base[i].numpy()[mask[i]] for i in range(pred.shape[0])]
+
+    figs = { type + "/garment_" + str(i): plot_3d(gt_curv[i], pred_curv[i], faces[i], mesh[i]) for i in range(4) }
+    
+    mean_hist = wandb.Histogram(sequence=mean.numpy().mean(axis=0), num_bins=model.latent_dim)
+    logvar_hist = wandb.Histogram(sequence=logvar.numpy().mean(axis=0), num_bins=model.latent_dim)
+
+    return { **figs, **{ type + '/images': imgs, type + '/mean_hist': mean_hist, type + '/logvar_hist': logvar_hist } }
+
+
+def load_model(model, checkpoint):
+    ckpt = tf.train.Checkpoint(model=model)
+    manager = tf.train.CheckpointManager(ckpt, checkpoint, max_to_keep=1)
+    ckpt.restore(manager.latest_checkpoint)
+
+
+def sample(folder):
+    load_model(model,folder)
+    for batch in dataset_train:
+        batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic))
+        uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
+        nr_of_samples = 16
+
+        # pred = model.sample(batch['uv_static'], batch['uv_body_posed'], batch['uv_mask'])
+        pred = model.sample_single(batch['uv_static'][0], batch['uv_body_posed'][0], batch['uv_mask'][0], nr_of_samples=nr_of_samples)
+
+        ce_uv_loss = scce(tf.tile(batch['uv_curvature'][0][tf.newaxis,:,:], [nr_of_samples,1,1]), pred)
+
+
+        mesh_pred = uv2mesh(pred).numpy().argmax(axis=-1)
+        mesh_orig = batch['mesh_curvature'][0].numpy()
+
+        F = np.load("custom_utils/smpl_faces.npy")
+        mask = batch['mesh_mask'].numpy().astype(bool)[0]
+        _, faces = _maskFaces(np.ones((14475, 3)), F, mask)
+        mesh_pred = [mesh_pred[i, mask] for i in range(mesh_pred.shape[0])] 
+        curv_color = [np.zeros((mesh_pred[i].shape[0], 3)) for i in range(len(mesh_pred))]
+        mesh_orig = mesh_orig[mask] + 1
+
+        colors = np.array([cm.color('jet', 1 - i) for i in np.linspace(0, 1, 6+1)])
+        curv_color_gt = np.zeros((mesh_pred[0].shape[0], 3))
+        for i in range(6+1):
+            for j in range(len(mesh_pred)):
+                curv_color[j][mesh_pred[j] == i, :] = colors[i]
+            curv_color_gt[mesh_orig == i, :] = colors[i]
+
+        mesh_base = batch['mesh_posed'][0].numpy()[mask, :]
+        meshes = []
+        meshes.append(go.Mesh3d(
+                    # 8 vertices of a cube
+                    x=mesh_base[:,0],
+                    y=mesh_base[:,1],
+                    z=mesh_base[:,2],
+                    # i, j and k give the vertices of triangles
+                    i = faces[:,0],
+                    j = faces[:,1],
+                    k = faces[:,2],
+                    vertexcolor=curv_color_gt,
+                    text='gt'
+                    ))
+
+        for i in range(nr_of_samples):
+            x = mesh_base.copy()
+            x[:, 1] += 1 * (i + 1)
+            meshes.append(go.Mesh3d(
+                        # 8 vertices of a cube
+                        x=x[:,0],
+                        y=x[:,1],
+                        z=x[:,2],
+                        # i, j and k give the vertices of triangles
+                        i = faces[:,0],
+                        j = faces[:,1],
+                        k = faces[:,2],
+                        vertexcolor=curv_color[i]
+                        ))
+
+
+
+
+        fig = go.Figure(data=meshes)
+
+        fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=nr_of_samples-1, z=1))
+        fig.update_layout(
+            scene = dict(
+                xaxis = dict(nticks=4, range=[-1,1],),
+                            yaxis = dict(nticks=4, range=[-1,nr_of_samples + 1],),
+                            zaxis = dict(nticks=4, range=[-1,1],),)
+                            )
+
+        # fig.show()
+
+        pred = pred.numpy().argmax(-1)
+        frames = [Image.fromarray(pred[i]/6*255) for i in range(pred.shape[0])]
+        frames[0].save("temp/curv2.gif", format="GIF", append_images=frames[1:],
+               save_all=True, duration=500, loop=0)
+
+        print()
+
+if __name__ == '__main__':
+    # init_wandb()
+    allow_memory_growth()
+    parser = argparse.ArgumentParser(description='')
+    parser.add_argument('--name', type=str)
+    args = vars(parser.parse_args())
+
+    batch_size = 4
+    
+    smpl = SMPL()
+    scce = tf.keras.losses.SparseCategoricalCrossentropy()
+
+    factory = DatasetFactory()
+    signature_dynamic = { 'id': (),
+                          'pose_shape': (72 + 11,),
+                        #   'uv_posed': (3, 512, 256,),
+                          'uv_body_posed': (3, 512, 256,),
+                          'mesh_body_posed': (14475, 3,),
+                          'mesh_posed': (14475, 3,),
+                        #   'mesh_unposed': (14475, 3,),
+                        #   'uv_unposed': (3, 512, 256,),
+                          'uv_curvature': (512, 256),
+                          'mesh_curvature':  (14475)
+                        }
+    signature_static = { 'uv_static': (3, 512, 256,),
+                        #  'mesh_static': (14475, 3,),
+                         'mesh_mask': (14475,),
+                         'uv_mask': (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)
+    static_dataset_train = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/train', signature=signature_static)
+    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',
+                 body_model='checpoints/ae_body_6',)
+
+
+    # lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(0.0001, decay_steps=30000, decay_rate=0.1, staircase=False)
+    optimizer = tf.optimizers.Adam(learning_rate=0.0001)
+    balancing_coeff = 0.1
+    reg_coeff = 1
+    folder = 'models/ae_logs/'
+    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)
+
+    step = 0
+    for batch in dataset_train:
+        batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic))
+
+        train_loss = train_step(batch)
+
+        # data_valid = valid_iterator.get_next()
+        # data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, signature_dynamic))
+        # valid_loss = valid_step(data_valid)
+        valid_loss = {}
+
+        to_log = { **train_loss, **valid_loss, **{ 'current_step': step} }
+        if step % 2500 == 0:
+            train_figs = plot_everything(batch, 'train')
+            # valid_figs = plot_everything(data_valid, 'valid')
+            valid_figs = {}
+            to_log = { **to_log, **train_figs, **valid_figs }
+
+        if step % 25000 == 0:
+            manager.save(checkpoint_number=step)
+
+        if step > 25000:
+        # if train_loss['train/euclid_error'] < 0.01:
+            balancing_coeff = balancing_coeff + 1 / 20000
+        # if train_loss['train/euclid_error'] < 0.015:
+        #     reg_coeff = reg_coeff + 100 / 15000
+
+
+        # if step > 15000 and balancing_coeff <= 1e-3:
+        #     balancing_coeff = 1e-2
+        # if step > 20000 and balancing_coeff <= 1e-3:
+        #     balancing_coeff = 1e-2
+        # if step > 25000 and balancing_coeff <= 1e-2:
+        #     balancing_coeff = 1e-1
+        # if step > 30000 and balancing_coeff <= 1e-1:
+        #     balancing_coeff = 1
+
+        # gamma = 0.99 * gamma + 0.01 * train_loss['train/reconstr_loss']        
+
+        to_log = { **to_log, **{ 'train/balancing_coeff': balancing_coeff, 'train/gamma': gamma } }
+        
+        wandb.log(to_log, step = step)
+
+        print(step)
+
+        step += 1

cvae_vae.py

@@ -0,0 +1,586 @@
+import os
+os.environ["CUDA_VISIBLE_DEVICES"] = "2"
+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 scipy as sp
+from custom_utils.custom_utils import UV2Mesh, maskFacesBatch, _maskFaces
+from tensorboardX import SummaryWriter
+import plotly.graph_objects as go
+import plotly.express as ex
+
+from utils import str_to_bool
+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 convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block
+from stylegan2_ae.generator import Generator
+from custom_utils.graph_utils import Reconstructor_V4_Decoder as GNNDecoder, sparse_to_tensor
+from custom_utils.graph_utils import Reconstructor_v2 as SimpleReconstructor
+from dataset import DatasetFactory, StaticDataset
+from custom_utils.custom_utils import batch_to_dict
+from custom_utils.smpl import SMPL
+from auto_encoder_body import AutoEncoder as BodyAE
+from auto_encoder_curv import AutoEncoder as CurvatureAE
+from auto_encoder_static import AutoEncoder as StaticAE
+from cvae_curv import CVAE as CurvatureVAE
+
+
+name='cvae_vae_1'
+def init_wandb():
+    run = wandb.init(project='vae', name=name,
+            notes='learn offsets. smaller flattened latent vectors. with 25-85k pretraining. smaller batchsize. fixed balancing of reconst and kl loss for baseline.\
+new implementation for kl loss. no lr schedule. 25 jumpsuits, average reconst_loss before calculation, linearly increasing balancing, up to 1000 from 250.\
+offsets over average unposed garment. gamma balancing. remove global orientation (in posed body input and in pose_shape parameters. \
+rebalanced l1 loss. retrained body encoder with deoriented bodies \
+')
+    artifact = wandb.Artifact('code', type='code')
+    artifact.add_file('cvae.py')
+    artifact.add_file('dataset.py')
+    artifact.add_file('losses.py')
+    artifact.add_file('convnext/models/convnext_tf.py')
+    run.log_artifact(artifact)
+
+uv2mesh = UV2Mesh(res=256)
+
+def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, collision_color):
+    mesh_pred_unposed[:,1] += 1
+    collision[:,1] += 2
+    fig = go.Figure(data=[
+        go.Mesh3d(
+            # 8 vertices of a cube
+            x=mesh_gt[:,0],
+            y=mesh_gt[:,1],
+            z=mesh_gt[:,2],
+            # i, j and k give the vertices of triangles
+            i = faces[:,0],
+            j = faces[:,1],
+            k = faces[:,2]
+        ),
+        go.Mesh3d(
+            # 8 vertices of a cube
+            x=mesh_pred_unposed[:,0],
+            y=mesh_pred_unposed[:,1],
+            z=mesh_pred_unposed[:,2],
+            # i, j and k give the vertices of triangles
+            i = faces[:,0],
+            j = faces[:,1],
+            k = faces[:,2],
+        ),
+        go.Mesh3d(
+            # 8 vertices of a cube
+            x=collision[:,0],
+            y=collision[:,1],
+            z=collision[:,2],
+            # i, j and k give the vertices of triangles
+            i = collision_faces[:,0],
+            j = collision_faces[:,1],
+            k = collision_faces[:,2],
+            vertexcolor=collision_color
+        )
+    ])
+
+    fig.update_layout(scene_aspectmode='manual',
+                    scene_aspectratio=dict(x=1, y=2, z=1))
+    fig.update_layout(
+        scene = dict(
+            xaxis = dict(nticks=4, range=[-1,1],),
+                        yaxis = dict(nticks=4, range=[-1,3],),
+                        zaxis = dict(nticks=4, range=[-1,1],),)
+                        )
+
+    return fig
+
+
+class Bottleneck(tf.keras.Model):
+    
+    def __init__(self, dim):
+        super(Bottleneck, self).__init__()
+
+        self.bottleneck = tf.keras.Sequential(
+                [Block(dim=dim, drop_path=0,
+                       layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
+            )
+
+    def call(self, inputs, training=None):
+        x = inputs
+
+        x = self.bottleneck(x, training=training)
+        
+        return x
+
+
+class Encoder(tf.keras.Model):
+    
+    def __init__(self, latent_dim):
+        super(Encoder, self).__init__()
+
+        self.encoder = ConvNeXt(depths=[3, 3, 9, 3], dims=[32, 64, 128, 256], include_top=False)
+        self.flatten = tf.keras.layers.Flatten()
+        self.mean = tf.keras.layers.Dense(latent_dim)
+        self.log_var = tf.keras.layers.Dense(latent_dim)
+
+    def call(self, inputs, training=None):
+        x = inputs
+
+        x = self.encoder(x)
+        x = self.flatten(x)
+        mean = self.mean(x)
+        log_var = self.log_var(x)
+
+        return mean, log_var
+
+
+class Decoder(tf.keras.Model):
+    
+    def __init__(self):
+        super(Decoder, self).__init__()
+
+        self.dense = tf.keras.layers.Dense(16 * 8 * 80)
+        self.reshape = tf.keras.layers.Reshape((16, 8, 80))
+        self.decoder = ConvNeXtReverse(depths=list(reversed([3, 3, 9, 3])), dims=list(reversed([256, 128, 64, 16])))
+
+
+    def call(self, inputs, training=None):
+        x = inputs
+
+        x = self.dense(x)
+        x = self.reshape(x)
+        x = self.decoder(x)
+
+        return x
+
+
+class CVAE(tf.keras.Model):
+    
+    def __init__(self, latent_dim=128):
+        super(CVAE, self).__init__()
+        self.latent_dim = latent_dim
+
+        model : StaticAE = StaticAE()
+        load_model(model, 'checkpoints/ae_static_3')
+        self.static_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
+        self.static_enc.trainable = False
+
+        model : BodyAE = BodyAE()
+        load_model(model, 'checpoints/ae_body_5')
+        self.body_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
+        self.body_enc.trainable = False
+
+        model : CurvatureAE = CurvatureVAE()
+        load_model(model, 'checpoints/cvae_curv_1_3')
+        self.curv_gen = model
+        self.curv_gen.trainable = False
+
+        self.encoder = Encoder(latent_dim)
+        self.decoder = Decoder()
+
+
+    # @tf.function
+    def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask, eps=None):
+        if eps is None:
+            eps = tf.random.normal(shape=(4, self.latent_dim))
+
+        uv_template = tf.tile(tf.transpose(uv_template, [1,2,0])[tf.newaxis,:,:,:], [4,1,1,1])
+        uv_body_posed = tf.tile(tf.transpose(uv_body_posed, [1,2,0])[tf.newaxis,:,:,:], [4,1,1,1])
+        uv_curv = tf.tile(uv_curv[tf.newaxis,:,:,:], [4,1,1,1])
+        uv_mask = tf.tile(uv_mask[tf.newaxis,:,:], [4,1,1])
+
+        # uv_template, uv_body_posed = tf.transpose(uv_template, [0,2,3,1]), 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_curv = self.curv_gen(uv_curv)
+
+        z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
+        z = tf.concat([eps, z_cond], axis=-1)
+
+        x = self.decoder(z)
+        x = x + avg_jumpsuit[tf.newaxis,:,:,:]
+        x = x * uv_mask[:,:,:,tf.newaxis]
+
+        return x
+
+
+    def reparameterize(self, mean, logvar):
+        eps = tf.random.normal(shape=mean.shape)
+        return eps * tf.exp(logvar * .5) + mean
+
+
+    def decode(self, z, apply_sigmoid=False):
+        logits = self.decoder(z)
+        if apply_sigmoid:
+            probs = tf.sigmoid(logits)
+            return probs
+        return logits
+
+
+    def call(self, inputs, training=None):
+        uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask, uv_shaped_body = inputs
+        uv_unposed, uv_template, uv_body_posed = tf.transpose(uv_unposed, [0,2,3,1]), tf.transpose(uv_template, [0,2,3,1]), 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_curv = self.curv_gen.sample(uv_template, uv_body_posed, uv_mask)
+
+        x = tf.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], axis=-1)
+
+        mean, logvar = self.encoder(x)
+        z = self.reparameterize(mean, logvar)
+        z = tf.concat([z, z_static, z_body, z_curv], axis=-1)
+        x = self.decoder(z)
+
+
+        # x = tf.transpose(x, [0,3,1,2])
+        # x = x * uv_mask[:,tf.newaxis,:,:]
+        # x = tf.image.resize(x, (512, 256), tf.image.ResizeMethod.LANCZOS5)
+        # x = x + tf.transpose(uv_shaped_body, [0,2,3,1])
+        x = x + avg_jumpsuit[tf.newaxis,:,:,:]
+        x = x * uv_mask[:,:,:,tf.newaxis]
+
+        return x, mean, logvar, z
+
+
+def log_normal_pdf(sample, mean, logvar, raxis=1):
+    log2pi = tf.math.log(2. * np.pi)
+    res = tf.reduce_sum(-.5 * ((sample - mean) ** 2. * tf.exp(-logvar) + logvar + log2pi), axis=raxis)
+    return res
+
+
+def kl_loss(mean, logvar):
+    kl_loss = 1 + logvar - tf.square(mean) - tf.exp(logvar)
+    kl_loss = tf.reduce_mean(kl_loss)
+    kl_loss *= -0.5
+    return kl_loss
+
+
+@tf.function
+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_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask'], batch['uv_shaped_body']])
+
+        mesh_pred = uv2mesh(tf.transpose(pred, [0,3,1,2]))
+        mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred)
+        mesh_orig = batch['mesh_unposed']
+
+
+        e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
+        c_loss = tf.reduce_mean(collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']))
+        l1_loss = tf.reduce_sum(tf.reduce_sum((tf.math.abs(tf.transpose(batch['uv_unposed'], [0,2,3,1]) - pred)), axis=-1) * batch['uv_mask'], axis=[1, 2])
+        # ms_ssim = tf.image.ssim_multiscale(tf.transpose(pred, [0,2,3,1]) + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=2)
+        ms_ssim = tf.image.ssim_multiscale(pred + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=3)
+        # ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale((tf.transpose(pred, [0,2,3,1]) + 1) * batch['uv_mask'][:,:,:,tf.newaxis], (tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1) * batch['uv_mask'][:,:,:,tf.newaxis], max_val=2))
+        rfft_loss = 0 #tf.reduce_mean(tf.reduce_sum(tf.math.abs(tf.signal.rfft3d(pred) - tf.signal.rfft3d(batch['uv_unposed'])), axis=[1,2,3]))
+        fft_loss = tf.reduce_sum(tf.math.abs(tf.signal.fft3d(tf.complex(pred, tf.zeros_like(pred))) - tf.signal.fft3d(tf.complex(tf.transpose(batch['uv_unposed'], [0,2,3,1]), tf.zeros_like(tf.transpose(batch['uv_unposed'], [0,2,3,1]))))), axis=[1,2,3])
+        fft_loss_pc, fft_loss_pc_low, fft_loss_pc_high = 0,0,0 #fft_loss_per_channel(pred, batch['uv_unposed'])
+        tv_loss = tf.image.total_variation(pred)
+
+        reconstr_loss = l1_loss / 4
+        reconstr_loss = tf.reduce_mean(reconstr_loss) + reg_coeff *(e_loss + c_loss)
+        # logpx_z = -reconstr_loss
+        # logpz = log_normal_pdf(z[:,:model.latent_dim], 0., 0.)
+        # logqz_x = log_normal_pdf(z[:,:model.latent_dim], mean, logvar)
+        # loss = -tf.reduce_mean(logpx_z + balancing_coeff * (logpz - logqz_x))
+        kl = kl_loss(mean, logvar)
+        loss = reconstr_loss / (2 * (gamma ** 2)) + balancing_coeff * kl
+        # loss = reconstr_loss + balancing_coeff * kl
+        ms_ssim = tf.reduce_mean(ms_ssim)
+        l1_loss = tf.reduce_mean(l1_loss)
+        fft_loss = tf.reduce_mean(fft_loss)
+        tv_loss = tf.reduce_mean(tv_loss)
+        mean = tf.reduce_mean(mean)
+        logvar = tf.reduce_mean(logvar)
+
+    mesh_orig_2 = uv2mesh(batch['uv_unposed'])
+    metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
+    metric_2 = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig_2 - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
+    euclid_posed = 0 # tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_skinned - batch['mesh_posed']) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
+    gradients = tape.gradient(loss, model.trainable_variables)
+    # gradients = [tf.clip_by_norm(g, 1.0) for g in gradients]
+    # gradients = [tf.clip_by_value(g, -1.0, 1.0) for g in gradients]
+    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
+    return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/collision_loss': c_loss, 'train/l1_loss': l1_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim,
+             'train/fft_loss': fft_loss, 'train/euclid_error_reconst': metric_2, 'train/rfft_loss': rfft_loss, 'train/fft_oc_low': fft_loss_pc_low, 'train/fft_pc_high': fft_loss_pc_high,
+             'train/fft_loss_per_channel': fft_loss_pc, 'train/euclid_error_posed': euclid_posed, 'train/reconstr_loss': reconstr_loss, 'train/tv_loss': tv_loss,
+             'train/mean': mean, 'train/logvar': logvar, 'train/kl_loss': kl, 'train/reconst_loss_norm': reconstr_loss / (2 * (gamma ** 2)), 'train/var': np.e ** logvar,
+             'train/reg_coeff': reg_coeff }
+        
+
+
+@tf.function
+def valid_step(batch):
+    uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
+    pred, mean, logvar, z = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask'], batch['uv_shaped_body']])
+
+    mesh_pred = uv2mesh(tf.transpose(pred, [0,3,1,2]))
+    mesh_orig = batch['mesh_unposed']
+    e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
+    c_loss = 1e-2 * tf.reduce_mean(collision_penalty(mesh_pred, batch['mesh_body_posed'], batch['mesh_mask']))
+    l1_loss = tf.reduce_sum(tf.reduce_sum((tf.math.abs(tf.transpose(batch['uv_unposed'], [0,2,3,1]) - pred)), axis=-1) * batch['uv_mask'], axis=[1, 2])
+    b_loss = 0 #1e-3 * bend_loss(mesh_pred, smpl_faces, laplacian_loss, batch['mesh_mask'])
+    # ms_ssim = tf.image.ssim_multiscale(tf.transpose(pred, [0,2,3,1]) + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=2)
+    ms_ssim = tf.image.ssim_multiscale(pred + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=3)
+    fft_loss = 0 #tf.reduce_sum(tf.math.abs(tf.signal.fft3d(tf.complex(pred, tf.zeros_like(pred))) - tf.signal.fft3d(tf.complex(tf.transpose(batch['uv_unposed'], [0,2,3,1]), tf.zeros_like(tf.transpose(batch['uv_unposed'], [0,2,3,1]))))), axis=[1,2,3])
+    # mesh_orig_2 = uv2mesh(batch['uv_unposed'])
+    metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
+    metric_2 = 0 #tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig_2 - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
+  
+    reconstr_loss = l1_loss
+    # logpx_z = -reconstr_loss
+    # logpz = log_normal_pdf(z[:,:model.latent_dim], 0., 0.)
+    # logqz_x = log_normal_pdf(z[:,:model.latent_dim], mean, logvar)
+    # loss = -tf.reduce_mean(logpx_z + balancing_coeff * (logpz - logqz_x))
+    kl = kl_loss(mean, logvar)
+    # loss = reconstr_loss / (2 * (gamma ** 2)) + kl
+    loss = reconstr_loss + balancing_coeff * kl
+    ms_ssim = tf.reduce_mean(ms_ssim)
+    l1_loss = tf.reduce_mean(l1_loss)
+    fft_loss = tf.reduce_mean(fft_loss)
+    b_loss = tf.reduce_mean(b_loss)
+    reconstr_loss = tf.reduce_mean(reconstr_loss)
+    mean = tf.reduce_mean(mean)
+    logvar = tf.reduce_mean(logvar)
+
+    return { 'valid/loss': loss, 'valid/edge_loss': e_loss, 'valid/collision_loss': c_loss, 'valid/l1_loss': l1_loss, 'valid/euclid_error': metric, 'valid/ms_ssim': ms_ssim,
+             'valid/fft_loss': fft_loss, 'valid/euclid_error_reconst': metric_2, 'train/bend_loss': b_loss, 'train/reconstr_loss': reconstr_loss,
+             'valid/mean': mean, 'valid/logvar': logvar, 'valid/kl_loss': kl, 'valid/var': np.e ** logvar }
+
+
+def plot_everything(batch: dict, type: str):
+    uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
+    pred, mean, logvar, z = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask'], batch['uv_shaped_body']])
+    
+    mesh_pred = uv2mesh(tf.transpose(pred, [0,3,1,2]))
+    mesh_orig = batch['mesh_unposed']
+    
+    orig_img = tf.transpose(batch['uv_unposed'], [0,2,3,1])
+    uv_static_img = tf.transpose(batch['uv_static'], [0,2,3,1])
+    uv_body_posed_img = tf.transpose(batch['uv_body_posed'], [0,2,3,1])
+    pred_img = pred
+    img = tf.concat((uv_static_img, uv_body_posed_img, orig_img, pred_img), axis=2)
+    imgs = [wandb.Image((img[i] + 1) / 2.0) for i in range(4)]
+
+    F = np.load("custom_utils/smpl_faces.npy")
+    mask = batch['mesh_mask'].numpy().astype(bool)
+    garment_unposed, faces = maskFacesBatch(mesh_pred.numpy(), F, mask)
+    gt = [mesh_orig[i].numpy()[mask[i], :] for i in range(pred.shape[0])]
+    
+    body = batch['mesh_body_posed'].numpy()
+    mesh_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred).numpy()
+    garment_posed = [mesh_posed[i, mask[i], :] for i in range(body.shape[0])]
+    collision = [np.concatenate((body[i], garment_posed[i]), axis=0) for i in range(len(garment_posed))]
+    collision_faces = [np.concatenate((F, faces[i] + 14475), axis=0) for i in range(len(garment_posed))]
+    collision_color = [np.concatenate((np.ones_like(body[i], dtype=np.int32)*100, np.ones_like(garment_posed[i], dtype=np.int32)*230), axis=0) for i in range(len(garment_posed))]
+
+    figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], faces[i], collision[i], collision_faces[i], collision_color[i]) for i in range(4) }
+    
+    mean_hist = wandb.Histogram(sequence=mean.numpy().mean(axis=0), num_bins=model.latent_dim)
+    logvar_hist = wandb.Histogram(sequence=logvar.numpy().mean(axis=0), num_bins=model.latent_dim)
+
+    return { **figs, **{ type + '/images': imgs, type + '/mean_hist': mean_hist, type + '/logvar_hist': logvar_hist } }
+
+
+def load_model(model, checkpoint):
+    ckpt = tf.train.Checkpoint(model=model)
+    manager = tf.train.CheckpointManager(ckpt, checkpoint, max_to_keep=1)
+    ckpt.restore(manager.latest_checkpoint)
+
+
+def sample(folder):
+    load_model(model,folder)
+    for batch in dataset_train:
+        batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic))
+        uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
+
+        pred = model.sample(batch['uv_static'][0], batch['uv_body_posed'][0], uv_curvature[0], batch['uv_mask'][0])
+        pose, shape = batch['pose_shape'][0, :72], batch['pose_shape'][0, 72:]
+        pose, shape = tf.tile(pose[tf.newaxis,:], [4,1]), tf.tile(shape[tf.newaxis,:], [4,1])
+
+        mesh_pred = uv2mesh(tf.transpose(pred, [0,3,1,2]))
+        mesh_posed = smpl.run_garment(pose, shape, mesh_pred).numpy()
+        mesh_orig = batch['mesh_posed'][0].numpy()
+
+        body = batch['mesh_body_posed'][0].numpy()
+        body_color = np.ones_like(body, dtype=np.int32)*100
+
+        F = np.load("custom_utils/smpl_faces.npy")
+        mask = batch['mesh_mask'].numpy().astype(bool)[0]
+        _, faces = _maskFaces(np.ones((14475, 3)), F, mask)
+
+        batch = 0
+
+        meshes = []
+        meshes.append(go.Mesh3d(
+                    # 8 vertices of a cube
+                    x=mesh_orig[mask,0],
+                    y=mesh_orig[mask,1],
+                    z=mesh_orig[mask,2],
+                    # i, j and k give the vertices of triangles
+                    i = faces[:,0],
+                    j = faces[:,1],
+                    k = faces[:,2]
+                    ))
+
+        for i in range(4):
+            x = mesh_posed[i][mask, :].copy()
+            x[:, 1] += 1 * (i + 1)
+            meshes.append(go.Mesh3d(
+                        # 8 vertices of a cube
+                        x=x[:,0],
+                        y=x[:,1],
+                        z=x[:,2],
+                        # i, j and k give the vertices of triangles
+                        i = faces[:,0],
+                        j = faces[:,1],
+                        k = faces[:,2]
+                        ))
+
+        for i in range(5):
+            x = body.copy()
+            x[:, 1] += 1 * i
+            meshes.append(go.Mesh3d(
+                        # 8 vertices of a cube
+                        x=x[:,0],
+                        y=x[:,1],
+                        z=x[:,2],
+                        # i, j and k give the vertices of triangles
+                        i = F[:,0],
+                        j = F[:,1],
+                        k = F[:,2],
+                        vertexcolor=body_color
+                        ))
+
+
+        fig = go.Figure(data=meshes)
+
+        fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=3, z=1))
+        fig.update_layout(
+            scene = dict(
+                xaxis = dict(nticks=4, range=[-1,1],),
+                            yaxis = dict(nticks=4, range=[-1,5],),
+                            zaxis = dict(nticks=4, range=[-1,1],),)
+                            )
+
+        fig.show()
+
+        print()
+
+if __name__ == '__main__':
+    # init_wandb()
+    # allow_memory_growth()
+    parser = argparse.ArgumentParser(description='')
+    parser.add_argument('--name', type=str)
+    args = vars(parser.parse_args())
+
+    batch_size = 12
+    
+    smpl = SMPL()
+
+    factory = DatasetFactory()
+    signature_dynamic = { 'id': (),
+                          'pose_shape': (72 + 11,),
+                        #   'uv_posed': (3, 512, 256,),
+                          'uv_body_posed': (3, 512, 256,),
+                          'mesh_body_posed': (14475, 3,),
+                          'mesh_posed': (14475, 3,),
+                          'mesh_unposed': (14475, 3,),
+                          'uv_unposed': (3, 512, 256,),
+                          'uv_curvature': (512, 256),
+                        }
+    signature_static = { 'uv_static': (3, 512, 256,),
+                        #  'mesh_static': (14475, 3,),
+                         'mesh_mask': (14475,),
+                         'uv_mask': (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)
+    static_dataset_train = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/train', signature=signature_static)
+    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()
+
+
+    # lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(0.0001, decay_steps=30000, decay_rate=0.1, staircase=False)
+    optimizer = tf.optimizers.Adam(learning_rate=0.0001)
+    balancing_coeff = 250
+    reg_coeff = 1
+    folder = 'models/ae_logs/'
+    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_10')
+
+
+    ckpt = tf.train.Checkpoint(model=model)
+    manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + name, max_to_keep=3)
+
+    step = 0
+    for batch in dataset_train:
+        batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic))
+
+        train_loss = train_step(batch)
+
+        # data_valid = valid_iterator.get_next()
+        # data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, signature_dynamic))
+        # valid_loss = valid_step(data_valid)
+        valid_loss = {}
+
+        to_log = { **train_loss, **valid_loss, **{ 'current_step': step} }
+        if step % 2500 == 0:
+            train_figs = plot_everything(batch, 'train')
+            # valid_figs = plot_everything(data_valid, 'valid')
+            valid_figs = {}
+            to_log = { **to_log, **train_figs, **valid_figs }
+
+        if step % 10000 == 0:
+            manager.save(checkpoint_number=step)
+
+        if train_loss['train/euclid_error'] < 0.01:
+            if balancing_activated_step == 0:
+                balancing_activated_step = step
+            balancing_coeff = balancing_coeff + 1000 / 30000
+        if train_loss['train/euclid_error'] < 0.015:
+            reg_coeff = reg_coeff + 100 / 15000
+        # if step > 15000 and balancing_coeff <= 1e-3:
+        #     balancing_coeff = 1e-2
+        # if step > 20000 and balancing_coeff <= 1e-3:
+        #     balancing_coeff = 1e-2
+        # if step > 25000 and balancing_coeff <= 1e-2:
+        #     balancing_coeff = 1e-1
+        # if step > 30000 and balancing_coeff <= 1e-1:
+        #     balancing_coeff = 1
+
+        # gamma = 0.99 * gamma + 0.01 * train_loss['train/reconstr_loss']        
+
+        to_log = { **to_log, **{ 'train/balancing_coeff': balancing_coeff, 'train/gamma': gamma } }
+        
+        wandb.log(to_log, step = step)
+
+        print(step)
+
+        step += 1

dataset.py

@@ -99,10 +99,11 @@ class StaticDataset():
 
 class DatasetIdx(tf.data.Dataset):
 
-    def __new__(cls, folder):
+    def __new__(cls, folder, shuffle=True):
         uv_static_folder = os.path.join(folder, 'pose_shape') 
-        frames = sorted(os.listdir(uv_static_folder))#[453:503]
+        frames = sorted(os.listdir(uv_static_folder))
         frames = [frame.split('.')[0] for frame in frames] # .split('_')[0]
+        if shuffle:
             random.shuffle(frames)
 
         return tf.data.Dataset.from_tensor_slices(frames)
@@ -126,9 +127,9 @@ class DatasetFactory():
         return tf_signature
 
     
-    def get_dataset(self, folder, batch_size, signature):
+    def get_dataset(self, folder, batch_size, signature, shuffle=True):
         # static_data = StaticDataset(folder, signature)
-        dataset_idx = DatasetIdx(folder)
+        dataset_idx = DatasetIdx(folder, shuffle)
 
         def read_item(frame):
             frame_split = tf.strings.split(frame, sep='_')
@@ -156,6 +157,9 @@ class DatasetFactory():
             # if 'mesh_body_posed' in signature.keys():
             #     mesh_body_posed = np.load(os.path.join(folder, 'mesh_body_posed', id + '_' + frame_nr + '.npy'))
             #     item['mesh_body_posed'] = mesh_body_posed
+            # if 'mesh_posed_body_no_orient' in signature.keys():
+            #     mesh_posed_body_no_orient = np.load(os.path.join(folder, 'mesh_posed_body_no_orient', id + '_' + frame_nr + '.npz')); mesh_posed_body_no_orient = mesh_posed_body_no_orient['mesh_posed_body_no_orient']
+            #     item['mesh_posed_body_no_orient'] = mesh_posed_body_no_orient
             if 'uv_body_posed' in signature.keys():
                 uv_posed_body_no_orient = np.load(os.path.join(folder, 'uv_posed_body_no_orient', id + '_' + frame_nr + '.npz')); uv_posed_body_no_orient = uv_posed_body_no_orient['uv_posed_body_no_orient']; uv_posed_body_no_orient = np.transpose(uv_posed_body_no_orient, [2, 0, 1])
                 item['uv_body_posed'] = uv_posed_body_no_orient
@@ -165,12 +169,18 @@ class DatasetFactory():
             if 'mesh_posed' in signature.keys():
                 mesh_posed = np.load(os.path.join(folder, 'mesh_posed', id + '_' + frame_nr + '.npz')); mesh_posed = mesh_posed['mesh_posed']
                 item['mesh_posed'] = mesh_posed
+            # if 'mesh_unposed' in signature.keys():
+            #     mesh_unposed = np.load(os.path.join(folder, 'mesh_unposed', id + '_' + frame_nr + '.npz')); mesh_unposed = mesh_unposed['mesh_unposed']
+            #     item['mesh_unposed'] = mesh_unposed
+            # if 'uv_unposed' in signature.keys():
+            #     uv_unposed = np.load(os.path.join(folder, 'uv_unposed', id + '_' + frame_nr + '.npz')); uv_unposed = uv_unposed['uv_unposed']; uv_unposed = np.transpose(uv_unposed, [2, 0, 1])
+            #     item['uv_unposed'] = uv_unposed
             if 'mesh_unposed' in signature.keys():
-                mesh_unposed = np.load(os.path.join(folder, 'mesh_unposed', id + '_' + frame_nr + '.npz')); mesh_unposed = mesh_unposed['mesh_unposed']
-                item['mesh_unposed'] = mesh_unposed
+                mesh_unposed_no_shape = np.load(os.path.join(folder, 'mesh_unposed_no_shape', id + '_' + frame_nr + '.npz')); mesh_unposed_no_shape = mesh_unposed_no_shape['mesh_unposed_no_shape']
+                item['mesh_unposed'] = mesh_unposed_no_shape
             if 'uv_unposed' in signature.keys():
-                uv_unposed = np.load(os.path.join(folder, 'uv_unposed', id + '_' + frame_nr + '.npz')); uv_unposed = uv_unposed['uv_unposed']; uv_unposed = np.transpose(uv_unposed, [2, 0, 1])
-                item['uv_unposed'] = uv_unposed
+                uv_unposed_no_shape = np.load(os.path.join(folder, 'uv_unposed_no_shape', id + '_' + frame_nr + '.npz')); uv_unposed_no_shape = uv_unposed_no_shape['uv_unposed_no_shape']; uv_unposed_no_shape = np.transpose(uv_unposed_no_shape, [2, 0, 1])
+                item['uv_unposed'] = uv_unposed_no_shape
             if 'uv_curvature' in signature.keys():
                 uv_curvature = np.load(os.path.join(folder, 'uv_curvature', id + '_' + frame_nr + '.npz')); uv_curvature = uv_curvature['curvature_uv']
                 item['uv_curvature'] = uv_curvature.astype(np.int32)

deshape.py

@@ -0,0 +1,117 @@
+import os
+import sys
+import time
+import cv2
+import random
+import numpy as np
+import tensorflow as tf
+import cmapy as cm
+import plotly.graph_objects as go
+import plotly.express as ex
+
+from dataset import DatasetFactory, StaticDataset
+from custom_utils.custom_utils import batch_to_dict, UV2Mesh, _maskFaces
+from custom_utils.smpl import SMPL
+
+if __name__ == '__main__':
+    os.environ["CUDA_VISIBLE_DEVICES"] = "5"
+
+    uv2mesh = UV2Mesh()
+    smpl = SMPL()
+
+    F = np.load("custom_utils/smpl_faces.npy")
+
+    factory = DatasetFactory()
+    signature_dynamic = { 'id': (),
+                          'pose_shape': (72 + 11,),
+                        #   'uv_posed': (3, 512, 256,),
+                        #   'uv_body_posed': (3, 512, 256,),
+                          # 'mesh_body_posed': (14475, 3,),
+                        #   'mesh_posed_body_no_orient': (14475, 3,),
+                          'mesh_posed': (14475, 3,),
+                          'mesh_unposed': (14475, 3,),
+                          # 'uv_unposed': (3, 512, 256,),
+                        #   'uv_curvature': (512, 256),
+                        #   'mesh_curvature':  (14475)
+                        }
+    signature_static = { #'uv_static': (3, 512, 256,),
+                        #  'mesh_static': (14475, 3,),
+                          'mesh_mask': (14475,),
+                        #  'uv_mask': (512, 256,),
+                        #  'uv_shaped_body': (3, 512, 256,),
+                          }
+
+    dataset = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/train',
+                                  batch_size=1,
+                                  signature=signature_dynamic)
+    static_dataset = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/train',
+                                   signature=signature_static)
+
+
+    uv_res = 256
+
+
+
+    for i, batch in enumerate(dataset):
+        print(i)
+        batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic))
+
+        mesh_unposed = np.squeeze(batch['mesh_unposed'].numpy())
+        mask = np.squeeze(batch['mesh_mask'].numpy()).astype(bool)
+        pose, shape = batch['pose_shape'][:, :72].numpy(), batch['pose_shape'][:, 72:].numpy()
+        # mesh_posed = np.squeeze(batch['mesh_posed'].numpy())
+        gender = int(shape[0, -1] > 0)
+        print(gender)
+
+        shape_offsets = tf.tensordot(shape[:,:10], smpl.shapedirs[gender], axes=[[1], [2]]).numpy()
+        
+
+        mesh_unposed[mask, :] = mesh_unposed[mask, :] - np.squeeze(shape_offsets)[smpl.to_keep, :][mask, :]
+
+        np.savez_compressed(os.path.join(folder, id, 'mesh_posed_body_no_orient', id + '_' + str(i) + '.npz'), mesh_posed_body_no_orient = T[i, :, :])
+
+        
+        # _, faces = _maskFaces(np.zeros((14475, 3), dtype=np.float32), F, mask)
+
+        # x = mesh_unposed[mask, :]
+        # y = mesh_unposed[mask, :] - np.squeeze(shape_offsets)[smpl.to_keep, :][mask, :]
+
+        # # x = smpl.v_template.numpy()[0, smpl.to_keep, :]
+        # # R = smpl.run_rodrigues(np.zeros((1, 72), dtype=np.float32))
+        # # x = smpl.run(R, shape)[0][0]
+        # # y = smpl.v_template.numpy()[gender, smpl.to_keep, :] + np.squeeze(shape_offsets)[smpl.to_keep, :]
+        # # faces = F
+
+        # fig = go.Figure(data=[
+        #     go.Mesh3d(
+        #         # 8 vertices of a cube
+        #         x=x[:,0],
+        #         y=x[:,1],
+        #         z=x[:,2],
+        #         # i, j and k give the vertices of triangles
+        #         i = faces[:,0],
+        #         j = faces[:,1],
+        #         k = faces[:,2]
+        #     ),
+        #     go.Mesh3d(
+        #         # 8 vertices of a cube
+        #         x=y[:,0],
+        #         y=y[:,1],
+        #         z=y[:,2],
+        #         # i, j and k give the vertices of triangles
+        #         i = faces[:,0],
+        #         j = faces[:,1],
+        #         k = faces[:,2]
+        #     )])
+
+        # fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1))
+        # fig.update_layout(
+        #     scene = dict(
+        #         xaxis = dict(nticks=4, range=[-1,1],),
+        #                     yaxis = dict(nticks=4, range=[-1,1],),
+        #                     zaxis = dict(nticks=4, range=[-1,1],),)
+        #                     )
+
+        # fig.show()
+
+        print()

models/cvae.py

@@ -0,0 +1,191 @@
+import os
+import time
+import random
+import tensorflow as tf
+import numpy as np
+import scipy as sp
+# import plotly.graph_objects as go
+# import plotly.express as ex
+
+from custom_utils.custom_utils import load_model
+from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block
+from auto_encoder_body import AutoEncoder as BodyAE
+from auto_encoder_curv import AutoEncoder as CurvatureAE
+from auto_encoder_static import AutoEncoder as StaticAE
+
+
+class Bottleneck(tf.keras.Model):
+    
+    def __init__(self, dim):
+        super(Bottleneck, self).__init__()
+
+        self.bottleneck = tf.keras.Sequential(
+                [Block(dim=dim, drop_path=0,
+                       layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
+            )
+
+    def call(self, inputs, training=None):
+        x = inputs
+
+        x = self.bottleneck(x, training=training)
+        
+        return x
+
+
+class Encoder(tf.keras.Model):
+    
+    def __init__(self, latent_dim):
+        super(Encoder, self).__init__()
+
+        self.encoder = ConvNeXt(depths=[3, 3, 9, 3], dims=[32, 64, 128, 256], include_top=False)
+        self.flatten = tf.keras.layers.Flatten()
+        self.mean = tf.keras.layers.Dense(latent_dim)
+        self.log_var = tf.keras.layers.Dense(latent_dim)
+
+    def call(self, inputs, training=None):
+        x = inputs
+
+        x = self.encoder(x)
+        x = self.flatten(x)
+        mean = self.mean(x)
+        log_var = self.log_var(x)
+
+        return mean, log_var
+
+
+class Decoder(tf.keras.Model):
+    
+    def __init__(self):
+        super(Decoder, self).__init__()
+
+        self.dense = tf.keras.layers.Dense(16 * 8 * 80)
+        self.reshape = tf.keras.layers.Reshape((16, 8, 80))
+        self.decoder = ConvNeXtReverse(depths=list(reversed([3, 3, 9, 3])), dims=list(reversed([256, 128, 64, 16])))
+
+
+    def call(self, inputs, training=None):
+        x = inputs
+
+        x = self.dense(x)
+        x = self.reshape(x)
+        x = self.decoder(x)
+
+        return x
+
+
+class CVAE(tf.keras.Model):
+    
+    def __init__(self, avg_garment_path='custom_data/jumpsuit_20_avg_uv.npy',
+                 latent_dim=128, static_model='checkpoints/ae_static_3',
+                 body_model='checpoints/ae_body_6',
+                 curvature_model='checpoints/ae_curvature_6'):
+        super(CVAE, self).__init__()
+        self.latent_dim = latent_dim
+
+        model : StaticAE = StaticAE()
+        load_model(model, static_model)
+        self.static_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
+        self.static_enc.trainable = False
+
+        model : BodyAE = BodyAE()
+        load_model(model, body_model)
+        self.body_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
+        self.body_enc.trainable = False
+
+        model : CurvatureAE = CurvatureAE()
+        load_model(model, curvature_model)
+        self.curv_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
+        self.curv_enc.trainable = False
+
+        self.encoder = Encoder(latent_dim)
+        self.decoder = Decoder()
+
+        self.avg_garment = np.load(avg_garment_path)
+
+        self.cond_encoder = tf.keras.Sequential([tf.keras.layers.Dense(512, activation='relu'),
+                                                 tf.keras.layers.Dense(256, activation='relu'),
+                                                 tf.keras.layers.Dense(128, activation='relu'),
+                                                 tf.keras.layers.Dense(latent_dim),])
+
+        # self.shuffle = tf.keras.layers.Lambda(lambda x: tf.stack([tf.gather(x[i], tf.random.shuffle(tf.range(x[i].shape[0]))) for i in range(x.shape[0])]))
+
+    # @tf.function
+    def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask, nr_of_samples=4, fix_noise=False, zero_latent=False):
+        if not fix_noise:
+            eps = tf.random.normal(shape=(nr_of_samples, self.latent_dim))
+        else:
+            eps = tf.random.normal(shape=(1, self.latent_dim))
+            eps = tf.tile(eps, [nr_of_samples,1])
+
+        uv_template = tf.tile(tf.transpose(uv_template, [1,2,0])[tf.newaxis,:,:,:], [nr_of_samples,1,1,1])
+        uv_body_posed = tf.tile(tf.transpose(uv_body_posed, [1,2,0])[tf.newaxis,:,:,:], [nr_of_samples,1,1,1])
+        if len(uv_curv.shape) == 3:
+            uv_curv = tf.tile(uv_curv[tf.newaxis,:,:,:], [nr_of_samples,1,1,1])
+        uv_mask = tf.tile(uv_mask[tf.newaxis,:,:], [nr_of_samples,1,1])
+
+        # uv_template, uv_body_posed = tf.transpose(uv_template, [0,2,3,1]), 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_curv = self.curv_enc(uv_curv)
+
+        z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
+        z_cond = self.cond_encoder(z_cond)
+
+        if zero_latent:
+            eps = tf.zeros_like(eps)
+
+        z = tf.concat([eps, z_cond], axis=-1)
+
+        x = self.decoder(z)
+        x = x + self.avg_garment[tf.newaxis,:,:,:]
+        x = x * uv_mask[:,:,:,tf.newaxis]
+
+        return x
+
+
+    def reparameterize(self, mean, logvar):
+        eps = tf.random.normal(shape=mean.shape)
+        return eps * tf.exp(logvar * .5) + mean
+
+
+    def decode(self, z, apply_sigmoid=False):
+        logits = self.decoder(z)
+        if apply_sigmoid:
+            probs = tf.sigmoid(logits)
+            return probs
+        return logits
+
+
+    def call(self, inputs, training=None, zero_latent=False):
+        uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask, uv_shaped_body = inputs
+        uv_unposed, uv_template, uv_body_posed = tf.transpose(uv_unposed, [0,2,3,1]), tf.transpose(uv_template, [0,2,3,1]), 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_curv = self.curv_enc(uv_curv)
+
+        x = tf.concat([uv_unposed, 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, z_curv], axis=-1)
+        z_cond = self.cond_encoder(z_cond)
+
+        if zero_latent:
+            z = tf.zeros_like(z)
+
+        z = tf.concat([z, z_cond], axis=-1)
+
+        # z = self.shuffle(z)
+
+        # z = tf.concat([z, z_curv], axis=-1)
+        x = self.decoder(z)
+
+
+        # x = tf.transpose(x, [0,3,1,2])
+        # x = x * uv_mask[:,tf.newaxis,:,:]
+        # x = tf.image.resize(x, (512, 256), tf.image.ResizeMethod.LANCZOS5)
+        # x = x + tf.transpose(uv_shaped_body, [0,2,3,1])
+        x = x + self.avg_garment[tf.newaxis,:,:,:]
+        x = x * uv_mask[:,:,:,tf.newaxis]
+
+        return x, mean, logvar, z

models/cvae_curv.py

@@ -0,0 +1,187 @@
+import os
+import time
+import random
+import tensorflow as tf
+import numpy as np
+import scipy as sp
+# import plotly.graph_objects as go
+# import plotly.express as ex
+
+from custom_utils.custom_utils import load_model
+from convnext.models.convnext_tf import ConvNeXt, Block, ConvNeXtReverseLabels
+from auto_encoder_body import AutoEncoder as BodyAE
+from auto_encoder_static import AutoEncoder as StaticAE
+
+
+
+class Bottleneck(tf.keras.Model):
+    
+    def __init__(self, dim):
+        super(Bottleneck, self).__init__()
+
+        self.bottleneck = tf.keras.Sequential(
+                [Block(dim=dim, drop_path=0,
+                       layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
+            )
+
+    def call(self, inputs, training=None):
+        x = inputs
+
+        x = self.bottleneck(x, training=training)
+        
+        return x
+
+
+class Encoder(tf.keras.Model):
+    
+    def __init__(self, latent_dim):
+        super(Encoder, self).__init__()
+
+        self.encoder = ConvNeXt(depths=[3, 3, 9, 3], dims=[32, 64, 128, 256], include_top=False)
+        self.flatten = tf.keras.layers.Flatten()
+        self.mean = tf.keras.layers.Dense(latent_dim)
+        self.log_var = tf.keras.layers.Dense(latent_dim)
+
+    def call(self, inputs, training=None):
+        x = inputs
+
+        x = self.encoder(x)
+        x = self.flatten(x)
+        mean = self.mean(x)
+        log_var = self.log_var(x)
+
+        return mean, log_var
+
+
+class Decoder(tf.keras.Model):
+    
+    def __init__(self):
+        super(Decoder, self).__init__()
+
+        self.dense = tf.keras.layers.Dense(16 * 8 * 80)
+        self.reshape = tf.keras.layers.Reshape((16, 8, 80))
+        self.decoder = ConvNeXtReverseLabels(depths=list(reversed([3, 3, 9, 3])), dims=list(reversed([256, 128, 64, 16])))
+
+
+    def call(self, inputs, training=None):
+        x = inputs
+
+        x = self.dense(x)
+        x = self.reshape(x)
+        x = self.decoder(x)
+
+        return x
+
+
+class CVAE(tf.keras.Model):
+    
+    def __init__(self, latent_dim=128, static_model='checkpoints/ae_static_3',
+                 body_model='checpoints/ae_body_6',):
+        super(CVAE, self).__init__()
+        self.latent_dim = latent_dim
+
+        model : StaticAE = StaticAE()
+        load_model(model, static_model)
+        self.static_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
+        self.static_enc.trainable = False
+
+        model : BodyAE = BodyAE()
+        load_model(model, body_model)
+        self.body_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
+        self.body_enc.trainable = False
+
+        self.cond_encoder = tf.keras.Sequential([tf.keras.layers.Dense(256, activation='relu'),
+                                                tf.keras.layers.Dense(256, activation='relu'),
+                                                tf.keras.layers.Dense(128, activation='relu'),
+                                                tf.keras.layers.Dense(latent_dim),])
+
+        self.encoder = Encoder(latent_dim)
+        self.decoder = Decoder()
+
+
+    # @tf.function
+    def sample_single(self, uv_template, uv_body_posed, uv_mask, eps=None, nr_of_samples=4):
+        if eps is None:
+            eps = tf.random.normal(shape=(nr_of_samples, self.latent_dim))
+
+        uv_template = tf.tile(tf.transpose(uv_template, [1,2,0])[tf.newaxis,:,:,:], [nr_of_samples,1,1,1])
+        uv_body_posed = tf.tile(tf.transpose(uv_body_posed, [1,2,0])[tf.newaxis,:,:,:], [nr_of_samples,1,1,1])
+        uv_mask = tf.tile(uv_mask[tf.newaxis,:,:], [nr_of_samples,1,1])
+
+        # uv_template, uv_body_posed = tf.transpose(uv_template, [0,2,3,1]), 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)
+        z = tf.concat([eps, z_cond], axis=-1)
+
+        x = self.decoder(z)
+        x = x * uv_mask[:,:,:,tf.newaxis]
+
+        return x
+
+
+    def sample(self, uv_template, uv_body_posed, uv_mask, eps=None):
+        eps = tf.random.normal(shape=(uv_template.shape[0], self.latent_dim))
+
+        uv_template = tf.transpose(uv_template, [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)
+        z = tf.concat([eps, z_cond], axis=-1)
+
+        x = self.decoder(z)
+        x = x * uv_mask[:,:,:,tf.newaxis]
+
+        return x
+
+
+    # encoders have to be the same for curv_vae and general CVAE
+    def sample_without_encoding(self, z_static, z_body, uv_mask, eps=None):
+        eps = tf.random.normal(shape=(z_static.shape[0], self.latent_dim))
+
+        z_cond = tf.concat([z_static, z_body], axis=-1)
+        z = tf.concat([eps, z_cond], axis=-1)
+
+        x = self.decoder(z)
+        x = x * uv_mask[:,:,:,tf.newaxis]
+
+        return x
+
+
+    def reparameterize(self, mean, logvar):
+        eps = tf.random.normal(shape=mean.shape)
+        return eps * tf.exp(logvar * .5) + mean
+
+
+    def decode(self, z, apply_sigmoid=False):
+        logits = self.decoder(z)
+        if apply_sigmoid:
+            probs = tf.sigmoid(logits)
+            return probs
+        return logits
+
+
+    def call(self, inputs, training=None):
+        uv_template, uv_body_posed, uv_curv, uv_mask, uv_shaped_body = inputs
+        uv_template, uv_body_posed = tf.transpose(uv_template, [0,2,3,1]), tf.transpose(uv_body_posed, [0,2,3,1])
+        z_static = self.static_enc(uv_template)
+        z_body = self.body_enc(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)
+
+        x = x * uv_mask[:,:,:,tf.newaxis]
+
+        return x, mean, logvar, z

sample.py

@@ -0,0 +1,332 @@
+import os
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+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 scipy as sp
+from custom_utils.custom_utils import UV2Mesh, maskFacesBatch, _maskFaces
+from tensorboardX import SummaryWriter
+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, 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 convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block
+from stylegan2_ae.generator import Generator
+from custom_utils.graph_utils import Reconstructor_V4_Decoder as GNNDecoder, sparse_to_tensor
+from custom_utils.graph_utils import Reconstructor_v2 as SimpleReconstructor
+from dataset import DatasetFactory, StaticDataset
+from custom_utils.custom_utils import batch_to_dict
+from custom_utils.smpl import SMPL
+from auto_encoder_body import AutoEncoder as BodyAE
+from auto_encoder_curv import AutoEncoder as CurvatureAE
+from auto_encoder_static import AutoEncoder as StaticAE
+from cvae_curv import CVAE as CVAECurv
+from models.cvae import CVAE as CVAEMain
+
+name='cvae_12_6'
+def init_wandb():
+    run = wandb.init(project='vae', name=name,
+            notes='learn offsets. smaller flattened latent vectors. with 25-85k pretraining. smaller batchsize. fixed balancing of reconst and kl loss for baseline. \
+new implementation for kl loss. no lr schedule. 25 jumpsuits, average reconst_loss before calculation, linearly increasing balancing, up to 1 from 0.1. \
+offsets over average unposed garment. remove global orientation (in posed body input and in pose_shape parameters. \
+retrained body encoder with deoriented bodies \
+use mean l2 loss. small batch. l2 * 1000, balancing to 1 from 0.1')
+    artifact = wandb.Artifact('code', type='code')
+    artifact.add_file('cvae.py')
+    artifact.add_file('dataset.py')
+    artifact.add_file('losses.py')
+    artifact.add_file('convnext/models/convnext_tf.py')
+    run.log_artifact(artifact)
+
+uv2mesh = UV2Mesh(res=256)
+
+def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, collision_color):
+    mesh_pred_unposed[:,2] += 1
+    collision[:,2] += 2
+    fig = go.Figure(data=[
+        go.Mesh3d(
+            # 8 vertices of a cube
+            x=mesh_gt[:,0],
+            y=mesh_gt[:,2],
+            z=mesh_gt[:,1],
+            # i, j and k give the vertices of triangles
+            i = faces[:,0],
+            j = faces[:,1],
+            k = faces[:,2]
+        ),
+        go.Mesh3d(
+            # 8 vertices of a cube
+            x=mesh_pred_unposed[:,0],
+            y=mesh_pred_unposed[:,2],
+            z=mesh_pred_unposed[:,1],
+            # i, j and k give the vertices of triangles
+            i = faces[:,0],
+            j = faces[:,1],
+            k = faces[:,2],
+        ),
+        go.Mesh3d(
+            # 8 vertices of a cube
+            x=collision[:,0],
+            y=collision[:,2],
+            z=collision[:,1],
+            # i, j and k give the vertices of triangles
+            i = collision_faces[:,0],
+            j = collision_faces[:,1],
+            k = collision_faces[:,2],
+            vertexcolor=collision_color
+        )
+    ])
+
+    fig.update_layout(scene_aspectmode='manual',
+                    scene_aspectratio=dict(x=1, y=2, z=1))
+    fig.update_layout(
+        scene = dict(
+            xaxis = dict(nticks=4, range=[-1,1],),
+                        yaxis = dict(nticks=4, range=[-1,3],),
+                        zaxis = dict(nticks=4, range=[-1,1],),)
+                        )
+
+    return fig
+
+
+@tf.function
+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_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask'], batch['uv_shaped_body']])
+
+
+
+def plot_everything(batch: dict, type: str):
+    uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
+    pred, mean, logvar, z = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask'], batch['uv_shaped_body']])
+    
+    mesh_pred = uv2mesh(tf.transpose(pred, [0,3,1,2]))
+    mesh_orig = batch['mesh_unposed']
+    
+    orig_img = tf.transpose(batch['uv_unposed'], [0,2,3,1])
+    uv_static_img = tf.transpose(batch['uv_static'], [0,2,3,1])
+    uv_body_posed_img = tf.transpose(batch['uv_body_posed'], [0,2,3,1])
+    pred_img = pred
+    img = tf.concat((uv_static_img, uv_body_posed_img, orig_img, pred_img), axis=2)
+    imgs = [wandb.Image((img[i] + 1) / 2.0) for i in range(4)]
+
+    F = np.load("custom_utils/smpl_faces.npy")
+    mask = batch['mesh_mask'].numpy().astype(bool)
+    garment_unposed, faces = maskFacesBatch(mesh_pred.numpy(), F, mask)
+    gt = [mesh_orig[i].numpy()[mask[i], :] for i in range(pred.shape[0])]
+    
+    body = batch['mesh_body_posed'].numpy()
+    mesh_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred).numpy()
+    garment_posed = [mesh_posed[i, mask[i], :] for i in range(body.shape[0])]
+    collision = [np.concatenate((body[i], garment_posed[i]), axis=0) for i in range(len(garment_posed))]
+    collision_faces = [np.concatenate((F, faces[i] + 14475), axis=0) for i in range(len(garment_posed))]
+    collision_color = [np.concatenate((np.ones_like(body[i], dtype=np.int32)*100, np.ones_like(garment_posed[i], dtype=np.int32)*230), axis=0) for i in range(len(garment_posed))]
+
+    figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], faces[i], collision[i], collision_faces[i], collision_color[i]) for i in range(4) }
+    
+    mean_hist = wandb.Histogram(sequence=mean.numpy().mean(axis=0), num_bins=model.latent_dim)
+    logvar_hist = wandb.Histogram(sequence=logvar.numpy().mean(axis=0), num_bins=model.latent_dim)
+
+    return { **figs, **{ type + '/images': imgs, type + '/mean_hist': mean_hist, type + '/logvar_hist': logvar_hist } }
+
+
+def load_model(model, checkpoint):
+    ckpt = tf.train.Checkpoint(model=model)
+    manager = tf.train.CheckpointManager(ckpt, checkpoint, max_to_keep=1)
+    ckpt.restore(manager.latest_checkpoint)
+
+
+def sample(folder_main, folder_curv):
+    # run = wandb.init(project='vae_plots', name=folder.split('/')[-1])
+    load_model(model, folder_main)
+    load_model(curv_model, folder_curv)
+    nr_of_samples = 2
+    count = 0
+    for batch in dataset_train:
+        batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic))
+
+        pred = curv_model.sample_single(batch['uv_static'][0], batch['uv_body_posed'][0], batch['uv_mask'][0], nr_of_samples=nr_of_samples)
+        uv_curvature = tf.math.argmax(pred, axis=-1)
+        uv_curvature = tf.one_hot(uv_curvature, 7)
+
+        pred = model.sample(batch['uv_static'][0], 
+                            batch['uv_body_posed'][0], 
+                            uv_curvature, 
+                            batch['uv_mask'][0], 
+                            nr_of_samples=nr_of_samples,
+                            fix_noise=True)
+        pose, shape = batch['pose_shape'][0, :72], batch['pose_shape'][0, 72:]
+        pose, shape = tf.tile(pose[tf.newaxis,:], [nr_of_samples,1]), tf.tile(shape[tf.newaxis,:], [nr_of_samples,1])
+
+        mesh_pred = uv2mesh(tf.transpose(pred, [0,3,1,2]))
+        mesh_posed = smpl.run_garment(pose, shape, mesh_pred).numpy()
+        mesh_orig = batch['mesh_posed'][0].numpy()
+        mesh_orig_reposed = smpl.run_garment(pose[:1], shape[:1], batch['mesh_unposed'][:1]).numpy()[0]
+
+        body = batch['mesh_body_posed'][0].numpy()
+        body_color = np.ones_like(body, dtype=np.int32)*100
+
+        F = np.load("custom_utils/smpl_faces.npy")
+        mask = batch['mesh_mask'].numpy().astype(bool)[0]
+        _, faces = _maskFaces(np.ones((14475, 3)), F, mask)
+
+
+        meshes = []
+        meshes.append(go.Mesh3d(
+                    # 8 vertices of a cube
+                    x=mesh_orig_reposed[mask,0],
+                    y=mesh_orig_reposed[mask,2],
+                    z=mesh_orig_reposed[mask,1],
+                    # i, j and k give the vertices of triangles
+                    i = faces[:,0],
+                    j = faces[:,1],
+                    k = faces[:,2]
+                    ))
+
+        for i in range(nr_of_samples):
+            x = mesh_posed[i][mask, :].copy()
+            # x[:, 2] += 1 * (i + 1)
+            x[:, 2] += 0.5
+            meshes.append(go.Mesh3d(
+                        # 8 vertices of a cube
+                        x=x[:,0],
+                        y=x[:,2],
+                        z=x[:,1],
+                        # i, j and k give the vertices of triangles
+                        i = faces[:,0],
+                        j = faces[:,1],
+                        k = faces[:,2]
+                        ))
+
+        # for i in range(5):
+        #     x = body.copy()
+        #     x[:, 2] += 1 * i
+        #     meshes.append(go.Mesh3d(
+        #                 # 8 vertices of a cube
+        #                 x=x[:,0],
+        #                 y=x[:,2],
+        #                 z=x[:,1],
+        #                 # i, j and k give the vertices of triangles
+        #                 i = F[:,0],
+        #                 j = F[:,1],
+        #                 k = F[:,2],
+        #                 vertexcolor=body_color
+        #                 ))
+
+
+        fig = go.Figure(data=meshes)
+
+        # fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=nr_of_samples-1, z=1))
+        fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1))
+        fig.update_layout(
+            scene = dict(
+                xaxis = dict(nticks=4, range=[-1,1],),
+                            # yaxis = dict(nticks=4, range=[-1,nr_of_samples+1],),
+                            yaxis = dict(nticks=4, range=[-1,1],),
+                            zaxis = dict(nticks=4, range=[-1,1],),)
+                            )
+
+        # wandb.log({'plot_' + str(count): fig})
+        fig.show()
+
+        uv_curvature_all = np.concatenate(uv_curvature.numpy().argmax(-1), axis=1)
+        ex.imshow(uv_curvature_all).show()
+
+        # frames = [Image.fromarray(pred[i].numpy()/6*255) for i in range(pred.shape[0])]
+        # frames[0].save("temp/curv_4.gif", format="GIF", append_images=frames[1:],
+        #        save_all=True, duration=500, loop=0, optimize=False)
+
+        if count > 10:
+            break
+
+        print(count)
+        count += 1
+
+        print()
+
+if __name__ == '__main__':
+    # init_wandb()
+    # allow_memory_growth()
+    parser = argparse.ArgumentParser(description='')
+    parser.add_argument('--name', type=str)
+    args = vars(parser.parse_args())
+
+    batch_size = 4
+    
+    smpl = SMPL()
+
+    factory = DatasetFactory()
+    signature_dynamic = { 'id': (),
+                          'pose_shape': (72 + 11,),
+                        #   'uv_posed': (3, 512, 256,),
+                          'uv_body_posed': (3, 512, 256,),
+                          'mesh_body_posed': (14475, 3,),
+                          'mesh_posed': (14475, 3,),
+                          'mesh_unposed': (14475, 3,),
+                          'uv_unposed': (3, 512, 256,),
+                          'uv_curvature': (512, 256),
+                        }
+    signature_static = { 'uv_static': (3, 512, 256,),
+                        #  'mesh_static': (14475, 3,),
+                         'mesh_mask': (14475,),
+                         'uv_mask': (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)
+    static_dataset_train = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/train', signature=signature_static)
+    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 = CVAEMain(avg_garment_path='custom_data/jumpsuit_20_avg_uv.npy',
+                    latent_dim=128, static_model='checkpoints/ae_static_3',
+                    body_model='checpoints/ae_body_6',
+                    curvature_model='checpoints/ae_curvature_6')
+    curv_model = CVAECurv()
+
+
+
+    # lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(0.0001, decay_steps=30000, decay_rate=0.1, staircase=False)
+    optimizer = tf.optimizers.Adam(learning_rate=0.0001)
+    balancing_coeff = 0.1
+    reg_coeff = 1
+    folder = 'models/ae_logs/'
+    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_13', 'checkpoints/cvae_curv_2')
+    # ckpt = tf.train.Checkpoint(model=model)
+    # manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + name, max_to_keep=3)
+
+    # step = 0
+    # for batch in dataset_train:
+    #     batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic))
+

verify_poses.py

@@ -0,0 +1,160 @@
+import os
+import sys
+import time
+import cv2
+import random
+import numpy as np
+import tensorflow as tf
+import cmapy as cm
+import plotly.graph_objects as go
+import plotly.express as ex
+
+from dataset import DatasetFactory, StaticDataset
+from custom_utils.custom_utils import batch_to_dict, UV2Mesh, _maskFaces
+from custom_utils.smpl import SMPL
+
+if __name__ == '__main__':
+    os.environ["CUDA_VISIBLE_DEVICES"] = "6"
+
+    uv2mesh = UV2Mesh()
+    smpl = SMPL()
+
+    F = np.load("custom_utils/smpl_faces.npy")
+
+    factory = DatasetFactory()
+    signature_dynamic = { 'id': (),
+                          'pose_shape': (72 + 11,),
+                        #   'uv_posed': (3, 512, 256,),
+                        #   'uv_body_posed': (3, 512, 256,),
+                          'mesh_body_posed': (14475, 3,),
+                        #   'mesh_posed_body_no_orient': (14475, 3,),
+                          'mesh_posed': (14475, 3,),
+                          'mesh_unposed': (14475, 3,),
+                          # 'uv_unposed': (3, 512, 256,),
+                        #   'uv_curvature': (512, 256),
+                        #   'mesh_curvature':  (14475)
+                        }
+    signature_static = { #'uv_static': (3, 512, 256,),
+                        #  'mesh_static': (14475, 3,),
+                         'mesh_mask': (14475,),
+                        #  'uv_mask': (512, 256,),
+                        #  'uv_shaped_body': (3, 512, 256,),
+                          }
+
+    dataset = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_test/train',
+                                  batch_size=1,
+                                  signature=signature_dynamic)
+    static_dataset = StaticDataset(folder='/home/workspace/local-data/jumpsuit_test/train',
+                                   signature=signature_static)
+
+
+    uv_res = 256
+
+
+
+    for i, batch in enumerate(dataset):
+        print(i)
+        batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic))
+
+        mesh_unposed = np.squeeze(batch['mesh_unposed'].numpy())
+        mask = np.squeeze(batch['mesh_mask'].numpy()).astype(bool)
+        pose, shape = batch['pose_shape'][:, :72].numpy(), batch['pose_shape'][:, 72:].numpy()
+        mesh_reposed = smpl.run_garment(pose, shape, batch['mesh_unposed']).numpy()[0]
+        mesh_posed = np.squeeze(batch['mesh_posed'].numpy())
+        mesh_body_posed = np.squeeze(batch['mesh_body_posed'].numpy())
+        mesh_posed_body_no_orient = np.squeeze(batch['mesh_posed_body_no_orient'].numpy())
+
+        # x = mesh_posed[mask, :]
+        R = smpl.run_rodrigues(pose)
+        # shape[:, -1] = 1
+        x = smpl.run(R, shape)[0][0]
+        y = mesh_body_posed
+        z = mesh_posed[mask, :]
+        _, faces = _maskFaces(np.zeros((14475, 3), dtype=np.float32), F, mask)
+        go.Figure(data=[
+                go.Mesh3d(
+                    # 8 vertices of a cube
+                    x=x[:,0],
+                    y=x[:,1],
+                    z=x[:,2],
+                    # i, j and k give the vertices of triangles
+                    i = F[:,0],
+                    j = F[:,1],
+                    k = F[:,2]
+                ),
+                go.Mesh3d(
+                    # 8 vertices of a cube
+                    x=y[:,0],
+                    y=y[:,1],
+                    z=y[:,2],
+                    # i, j and k give the vertices of triangles
+                    i = F[:,0],
+                    j = F[:,1],
+                    k = F[:,2]
+                ),
+                # go.Mesh3d(
+                #     # 8 vertices of a cube
+                #     x=z[:,0],
+                #     y=z[:,1],
+                #     z=z[:,2],
+                #     # i, j and k give the vertices of triangles
+                #     i = faces[:,0],
+                #     j = faces[:,1],
+                #     k = faces[:,2]
+                # ),
+                ]).show()
+
+        pose_2 = pose.copy()
+        pose_2[:, :3] = 0
+        R = smpl.run_rodrigues(pose_2)
+        x = smpl.run(R, shape)[0][0]
+        y = mesh_posed_body_no_orient
+        _, faces = _maskFaces(np.zeros((14475, 3), dtype=np.float32), F, mask)
+        go.Figure(data=[
+                go.Mesh3d(
+                    # 8 vertices of a cube
+                    x=x[:,0],
+                    y=x[:,1],
+                    z=x[:,2],
+                    # i, j and k give the vertices of triangles
+                    i = F[:,0],
+                    j = F[:,1],
+                    k = F[:,2]
+                ),
+                go.Mesh3d(
+                    # 8 vertices of a cube
+                    x=y[:,0],
+                    y=y[:,1],
+                    z=y[:,2],
+                    # i, j and k give the vertices of triangles
+                    i = F[:,0],
+                    j = F[:,1],
+                    k = F[:,2]
+                )]).show()
+
+        x = mesh_posed[mask, :]
+        y = mesh_reposed[mask, :]
+        _, faces = _maskFaces(np.zeros((14475, 3), dtype=np.float32), F, mask)
+        go.Figure(data=[
+                go.Mesh3d(
+                    # 8 vertices of a cube
+                    x=x[:,0],
+                    y=x[:,1],
+                    z=x[:,2],
+                    # i, j and k give the vertices of triangles
+                    i = faces[:,0],
+                    j = faces[:,1],
+                    k = faces[:,2]
+                ),
+                go.Mesh3d(
+                    # 8 vertices of a cube
+                    x=y[:,0],
+                    y=y[:,1],
+                    z=y[:,2],
+                    # i, j and k give the vertices of triangles
+                    i = faces[:,0],
+                    j = faces[:,1],
+                    k = faces[:,2]
+                )]).show()
+
+        print()

verify_scales.py

@@ -8,6 +8,7 @@ import tensorflow as tf
 
 from dataset import DatasetFactory, StaticDataset
 from custom_utils.custom_utils import batch_to_dict, UV2Mesh, _maskFaces
+from custom_utils.smpl import SMPL
 import plotly.graph_objects as go
 import plotly.express as ex
 import cmapy as cm
@@ -16,57 +17,81 @@ if __name__ == '__main__':
     os.environ["CUDA_VISIBLE_DEVICES"] = "6"
 
     uv2mesh = UV2Mesh()
+    smpl = SMPL()
 
     F = np.load("custom_utils/smpl_faces.npy")
 
     factory = DatasetFactory()
     signature_dynamic = { 'id': (),
-                        #   'pose_shape': (72 + 11,),
+                          'pose_shape': (72 + 11,),
                         #   'uv_posed': (3, 512, 256,),
                         #   'uv_body_posed': (3, 512, 256,),
                         #   'mesh_body_posed': (14475, 3,),
                         #   'mesh_posed': (14475, 3,),
                           'mesh_unposed': (14475, 3,),
-                          'uv_unposed': (3, 512, 256,),
+                        #   'uv_unposed': (3, 512, 256,),
                         #   'uv_curvature': (512, 256),
                         #   'mesh_curvature':  (14475)
                         }
     signature_static = { #'uv_static': (3, 512, 256,),
                         #  'mesh_static': (14475, 3,),
                          'mesh_mask': (14475,),
-                         'uv_mask': (512, 256,),
+                        #  'uv_mask': (512, 256,),
                         #  'uv_shaped_body': (3, 512, 256,),
                           }
 
     dataset = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/train',
                                   batch_size=1,
-                                  signature=signature_dynamic)
+                                  signature=signature_dynamic,
+                                  shuffle=False)
     static_dataset = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/train',
                                    signature=signature_static)
 
-
+    meshes = []
+    current_id = ''
+    count = 0
     for i, batch in enumerate(dataset):
+        id = batch[0].numpy()[0].decode('utf-8')
+        if id == current_id:
+            continue
+        current_id = id
+
         print(i)
         batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic))
 
-        uv_unposed = np.squeeze(batch['uv_unposed'].numpy())
+        # uv_unposed = np.squeeze(batch['uv_unposed'].numpy())
         mesh_unposed = np.squeeze(batch['mesh_unposed'].numpy())
         mask = np.squeeze(batch['mesh_mask'].numpy()).astype(bool)
-        _, faces = _maskFaces(np.ones((14475, 3)), F, mask)
+        _, f = _maskFaces(np.ones((14475, 3)), F, mask)
 
-        mesh_unposed[:,1] += (1 * i)
-        fig = go.Figure(data=[
-            go.Mesh3d(
+        # deshaped one
+        ###
+        pose, shape = batch['pose_shape'][:, :72].numpy(), batch['pose_shape'][:, 72:].numpy()
+        gender = int(shape[0, -1] > 0)
+        if gender == 0:
+            continue
+
+        shape_offsets = tf.tensordot(shape[:,:10], smpl.shapedirs[gender], axes=[[1], [2]]).numpy()
+        mesh_unposed -= np.squeeze(shape_offsets)[smpl.to_keep, :]
+        ###
+
+        mesh_unposed[:,2] += (0.5 * count)
+        meshes.append(go.Mesh3d(
                 # 8 vertices of a cube
                 x=mesh_unposed[mask,0],
                 y=mesh_unposed[mask,2],
                 z=mesh_unposed[mask,1],
                 # i, j and k give the vertices of triangles
-                i = faces[:,0],
-                j = faces[:,1],
-                k = faces[:,2],
-            )
-        ])
+                i = f[:,0],
+                j = f[:,1],
+                k = f[:,2],
+            ))
+
+        count += 1
+        if count > 10:
+            break
+
+    fig = go.Figure(data=meshes)
 
     fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=10, z=1))
     fig.update_layout(