pro sampling avaiable, updated cvae

cvae.py

@@ -1,5 +1,5 @@
 import os
-os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+os.environ["CUDA_VISIBLE_DEVICES"] = "2"
 import tensorflow as tf
 
 gpus = tf.config.experimental.list_physical_devices('GPU')
@@ -21,7 +21,6 @@ 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
@@ -38,7 +37,7 @@ from custom_utils.smpl import SMPL
 from models.cvae_model import CVAE
 
 
-name='cvae_16'
+name='cvae_17'
 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. \
@@ -46,7 +45,7 @@ new implementation for kl loss. no lr schedule. 25 jumpsuits, average reconst_lo
 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\
+fixed gender issue and retrained body encoder, FINAL\
 remove shape from unposed garment\
 training on 100 garments with offsets over the average of these')
     artifact = wandb.Artifact('code', type='code')
@@ -129,35 +128,29 @@ 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']])
+        pred, mean, logvar, z = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask']], training=True, use_mask=True)
 
-        mesh_pred = uv2mesh(tf.transpose(pred, [0,3,1,2]))
+        mesh_pred = uv2mesh(pred)
-        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]
+        mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
+        mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)
 
         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.image.ssim_multiscale(pred + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=3, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1))
         # 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 = 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])
-        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 = 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 + balancing_coeff * kl
         # loss = reconstr_loss + balancing_coeff * kl
@@ -178,8 +171,8 @@ def train_step(batch):
     # 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': fft_loss, 'train/euclid_error_reconst': metric_2,
-             '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/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/l2_loss': l2_loss }
         
@@ -226,9 +219,9 @@ def valid_step(batch):
 
 def plot_everything(batch: dict, type: str):
     uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
-    pred, mean, logvar, z = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask'], batch['uv_shaped_body']])
+    pred, mean, logvar, z = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask']], training=False, use_mask=True)
     
-    mesh_pred = uv2mesh(tf.transpose(pred, [0,3,1,2]))
+    mesh_pred = uv2mesh(pred).numpy()
     mesh_orig = batch['mesh_unposed']
     
     orig_img = tf.transpose(batch['uv_unposed'], [0,2,3,1])
@@ -237,20 +230,23 @@ def plot_everything(batch: dict, type: str):
     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()
-
     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])]
+    F = np.load("custom_utils/smpl_faces.npy")
+    mask = batch['mesh_mask'].numpy().astype(bool)
+    garment_unposed, faces = maskFacesBatch(mesh_pred, F, mask)
+    garment_unposed = [garment_unposed[i] + shape_offsets.numpy()[i, mask[i], :] for i in range(mask.shape[0])]
+    gt = [mesh_orig[i].numpy()[mask[i], :] + shape_offsets.numpy()[i, mask[i], :] for i in range(pred.shape[0])]
+    
+    body = batch['mesh_body_posed'].numpy()
+    for i in range(body.shape[0]):
+        mesh_pred[i, mask[i], :] += shape_offsets.numpy()[i, mask[i], :]
+    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))]
@@ -274,7 +270,7 @@ def sample(folder):
         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], nr_of_samples=nr_of_samples)
+        pred = model.sample_single(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,:], [nr_of_samples,1]), tf.tile(shape[tf.newaxis,:], [nr_of_samples,1])
 
@@ -291,7 +287,7 @@ def sample(folder):
         _, faces = _maskFaces(np.ones((14475, 3)), F, mask)
 
 
-        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)
+        pred_zero = model.sample_single(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()
 
@@ -412,18 +408,18 @@ if __name__ == '__main__':
                         #  'mesh_static': (14475, 3,),
                          'mesh_mask': (14475,),
                          'uv_mask': (512, 256,),
-                         'uv_shaped_body': (3, 512, 256,),
+                        #  'uv_shaped_body': (3, 512, 256,),
                         }
 
-    dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100', batch_size=batch_size, signature=signature_dynamic)
+    dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100', batch_size=batch_size, signature=signature_dynamic, use_filled=False)
-    static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', signature=signature_static)
+    static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', signature=signature_static, use_filled=False)
     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(avg_garment_path='custom_data/train_100_no_shape_avg_uv.npy',
                  latent_dim=128, static_model='checkpoints/ae_static_4',
-                 body_model='checpoints/ae_body_7',
+                 body_model='checpoints/ae_body_8',
                  curvature_model='checpoints/ae_curvature_7')
 
 

models/cvae_model.py

@@ -126,7 +126,7 @@ class CVAE(tf.keras.Model):
         # 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):
+    def sample_single(self, uv_template, uv_body_posed, uv_curv, uv_mask, nr_of_samples=4, fix_noise=False):
         if not fix_noise:
             eps = tf.random.normal(shape=(nr_of_samples, self.latent_dim))
         else:
@@ -147,9 +147,6 @@ class CVAE(tf.keras.Model):
         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)
@@ -159,6 +156,26 @@ class CVAE(tf.keras.Model):
         return x
 
 
+    def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask=None, noise=None):
+        if noise is None:
+            eps = tf.random.normal(shape=(uv_template.shape[0], self.latent_dim))
+        else:
+            eps = noise
+
+        z_cond = self.encode_conds(uv_template, uv_body_posed, uv_curv)
+        z = tf.concat([eps, z_cond], axis=-1)
+
+        x = self.decoder(z)
+        
+        if self.avg_garment is not None:
+            x = x + self.avg_garment[tf.newaxis,:,:,:]
+        
+        if uv_mask is not None:
+            x = x * uv_mask[:,:,:,tf.newaxis]
+
+        return x
+
+
     def reparameterize(self, mean, logvar, fix_noise=False):
         if not fix_noise:
             eps = tf.random.normal(shape=mean.shape)
@@ -176,22 +193,37 @@ class CVAE(tf.keras.Model):
         return logits
 
 
-    def call(self, inputs, training=None, zero_latent=False, fix_noise=False):
+    def encode_conds(self, uv_template, uv_body_posed, uv_curv):
-        uv_unposed, uv_template, uv_body_posed, uv_curv = inputs
+        if uv_template.shape[1] == 3:
-        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])
+            uv_template = tf.transpose(uv_template, [0,2,3,1])
+        if uv_body_posed.shape[1] == 3:
+            uv_body_posed = tf.transpose(uv_body_posed, [0,2,3,1])
+        if uv_curv.shape[1] == 3:
+            uv_curv = tf.transpose(uv_curv, [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, fix_noise)
         z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
         z_cond = self.cond_encoder(z_cond)
 
-        if zero_latent:
+        return z_cond
-            z = tf.zeros_like(z)
+
+
+    def call(self, inputs, training=None, use_mask=False):
+        if use_mask:
+            uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask = inputs
+        else:
+            uv_unposed, uv_template, uv_body_posed, uv_curv = 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])
+
+        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 = self.encode_conds(uv_template, uv_body_posed, uv_curv)
 
         z = tf.concat([z, z_cond], axis=-1)
 
@@ -200,14 +232,14 @@ class CVAE(tf.keras.Model):
         # 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])
         if self.avg_garment is not None:
             x = x + self.avg_garment[tf.newaxis,:,:,:]
-        # x = x * uv_mask[:,:,:,tf.newaxis]
+        if use_mask:
+            x = x * uv_mask[:,:,:,tf.newaxis]
 
         return x, mean, logvar, z
 

sample.py

@@ -174,7 +174,7 @@ def sample(folder_main, folder_curv):
         uv_curvature = tf.math.argmax(pred, axis=-1)
         uv_curvature = tf.one_hot(uv_curvature, 7)
 
-        pred = model.sample(batch['uv_static'][0], 
+        pred = model.sample_single(batch['uv_static'][0], 
                             batch['uv_body_posed'][0], 
                             uv_curvature, 
                             batch['uv_mask'][0], 

sample_pro.py

@@ -1,4 +1,3 @@
-from lib2to3.refactor import get_fixers_from_package
 import os
 os.environ["CUDA_VISIBLE_DEVICES"] = "0"
 import tensorflow as tf
@@ -22,27 +21,12 @@ 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 plotly.subplots import make_subplots
 
-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_model import CVAE as CVAEMain
 
@@ -163,84 +147,251 @@ def get_frame(folder, nr):
     frames = sorted(os.listdir(pose_folder))
     frames = [frame.split('.')[0] for frame in frames] # .split('_')[0]
     random.shuffle(frames)
-    return frames[0], frames[1:nr]
+    out = [frames[0]]
+    out.extend(frames[1:nr])
+    return out
 
-def get_bodies(folder, frame_main, frame_rand, nr=4, random=True):
+
+def get_bodies(folder, frames, random=True):
     bodies = []
-    uv_posed_body_no_orient = np.load(os.path.join(folder, 'uv_posed_body_no_orient', frame_main + '.npz')); uv_posed_body_no_orient = uv_posed_body_no_orient['uv_posed_body_no_orient']; bodies.append(uv_posed_body_no_orient)
+    poses_shapes = []
-    pass
+    uv_posed_body_no_orient = np.load(os.path.join(folder, 'uv_posed_body_no_orient', frames[0] + '.npz')); uv_posed_body_no_orient = uv_posed_body_no_orient['uv_posed_body_no_orient']; bodies.append(uv_posed_body_no_orient)
-
+    pose_shape = np.load(os.path.join(folder, 'pose_shape', frames[0] + '.npy')); poses_shapes.append(pose_shape)
-def get_curvs(folder, body, static, mask, frame=None, nr=4, random=True, model=None):
+    if not random:
-    assert (random and model is not None) or (not random and model is None)
+        for _ in range(len(frames)-1):
-    if random:
+            bodies.append(uv_posed_body_no_orient)
-        curv = model([body, static, mask])
+            poses_shapes.append(pose_shape)
-
-
-        uv_curvature = tf.math.argmax(pred, axis=-1)
-        uv_curvature = tf.one_hot(uv_curvature, 7)
-        pass
     else:
-        curv = np.load(os.path.join(folder, 'uv_curvature')) #...
+        for i in range(1, len(frames)):
-        pass    
+            uv_posed_body_no_orient = np.load(os.path.join(folder, 'uv_posed_body_no_orient', frames[i] + '.npz')); uv_posed_body_no_orient = uv_posed_body_no_orient['uv_posed_body_no_orient']; bodies.append(uv_posed_body_no_orient)
+            pose_shape = np.load(os.path.join(folder, 'pose_shape', frames[i] + '.npy')); poses_shapes.append(pose_shape)
+    bodies = np.stack(bodies)
+    poses_shapes = np.stack(poses_shapes).astype(np.float32)
+    return bodies, poses_shapes
 
-def get_statics(folder, nr=4, random=True):
-    pass
 
-def get_unposed(folder, nr=4, random=True):
+def get_curvs(folder, body, static, mask, frames, noise, random=True, model: CVAECurv=None, use_original=True):
-    pass
+    assert (random and model is not None) or (not random)
+    uv_curvature_orig = np.load(os.path.join(folder, 'uv_curvature', frames[0] + '.npz')); uv_curvature_orig = uv_curvature_orig['curvature_uv']
 
-def plot_inputs(unposed, static, body, curv):
+    noise = get_noise(not random, body.shape[0], latent_dim=128)
-    pass
+    curv = model.sample(static, body, mask, noise)
 
-def plot_variable_condition(condition, is_curv=False):
+    uv_curvature = tf.math.argmax(curv, axis=-1).numpy()
-    pass
+    if use_original:
+        uv_curvature[0,:,:] = uv_curvature_orig
+    out = tf.one_hot(uv_curvature, 7).numpy()
 
-def plot_results(unposed, posed, body, mesh_mask):
+    return out
-    body_color = np.ones_like(body, dtype=np.int32)*100
 
+def get_statics(folder, frames, random=True):
+    statics = []
+    uv_masks = []
+    mesh_masks = []
+    uv_static = np.load(os.path.join(folder, 'uv_static', frames[0].split('_')[0] +  '.npz')); uv_static = uv_static['uv_static']; statics.append(uv_static)
+    uv_mask = np.load(os.path.join(folder, 'uv_mask', frames[0].split('_')[0] + '.npz')); uv_mask=uv_mask['uv_mask']; uv_mask = uv_mask.astype(np.float32); uv_masks.append(uv_mask)
+    mesh_mask = np.load(os.path.join(folder, 'mesh_mask', frames[0].split('_')[0] + '.npy')); mesh_masks.append(mesh_mask)
+
+    if not random:
+        for _ in range(len(frames)-1):
+            statics.append(uv_static)
+            uv_masks.append(uv_mask)
+            mesh_masks.append(mesh_mask)
+    else:
+        for i in range(1, len(frames)):
+            uv_static = np.load(os.path.join(folder, 'uv_static', frames[i].split('_')[0] +  '.npz')); uv_static = uv_static['uv_static']; statics.append(uv_static)
+            uv_mask = np.load(os.path.join(folder, 'uv_mask', frames[i].split('_')[0] + '.npz')); uv_mask=uv_mask['uv_mask']; uv_mask = uv_mask.astype(np.float32); uv_masks.append(uv_mask)
+            mesh_mask = np.load(os.path.join(folder, 'mesh_mask', frames[i].split('_')[0] + '.npy')); mesh_masks.append(mesh_mask)
+    
+    statics = np.stack(statics)
+    uv_masks = np.stack(uv_masks)
+    mesh_masks = np.stack(mesh_masks)
+    return statics, uv_masks, mesh_masks
+
+
+def get_unposed(folder, frames, random=True):
+    uv_unposeds = []
+    mesh_unposeds = []
+    uv_unposed = np.load(os.path.join(folder, 'uv_posed', frames[0] + '.npz')); uv_unposed = uv_unposed['uv_gt']; uv_unposeds.append(uv_unposed)
+    mesh_unposed = np.load(os.path.join(folder, 'mesh_unposed_no_shape', frames[0] + '.npz')); mesh_unposed = mesh_unposed['mesh_unposed_no_shape']; mesh_unposeds.append(mesh_unposed)
+    if not random:
+        for _ in range(len(frames)-1):
+            uv_unposeds.append(uv_unposed)
+            mesh_unposeds.append(mesh_unposed)
+    else:
+        for i in range(1, len(frames)):
+            uv_static = np.load(os.path.join(folder, 'uv_static', frames[i].split('_')[0] +  '.npz')); uv_static = uv_static['uv_static']; uv_unposeds.append(uv_static)
+            mesh_unposed = np.load(os.path.join(folder, 'mesh_unposed_no_shape', frames[i] + '.npz')); mesh_unposed = mesh_unposed['mesh_unposed_no_shape']; mesh_unposeds.append(mesh_unposed)
+    uv_unposeds = np.stack(uv_unposeds)
+    mesh_unposeds = np.stack(mesh_unposeds)
+    return uv_unposeds, mesh_unposeds
+
+
+def make_square_2(figure):
+    figure.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1),
+                         scene2_aspectmode='manual', scene2_aspectratio=dict(x=1, y=1, z=1))
+    helper = dict(xaxis = dict(nticks=4, range=[-1,1],),
+                  yaxis = dict(nticks=4, range=[-1,1],),
+                  zaxis = dict(nticks=4, range=[-1,1],),)
+    figure.update_layout(scene = helper,
+                         scene2 = helper)
+    return figure
+
+def make_square_4(figure):
+    figure.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1),
+                         scene2_aspectmode='manual', scene2_aspectratio=dict(x=1, y=1, z=1),
+                         scene3_aspectmode='manual', scene3_aspectratio=dict(x=1, y=1, z=1),
+                         scene4_aspectmode='manual', scene4_aspectratio=dict(x=1, y=1, z=1))
+    helper = dict(xaxis = dict(nticks=4, range=[-1,1],),
+                  yaxis = dict(nticks=4, range=[-1,1],),
+                  zaxis = dict(nticks=4, range=[-1,1],),)
+    figure.update_layout(scene = helper,
+                         scene2 = helper,
+                         scene3 = helper,
+                         scene4 = helper,)
+    return figure
+
+
+def plot_inputs(unposed, static, body, curv, mesh_mask):
+    curv = np.tile(curv.argmax(-1)[:,:,np.newaxis], [1,1,3])/6
+    img = np.concatenate([unposed, static, body, curv], axis=-2)
+    ex.imshow((img + 1) / 2).show()
+
+    mesh_static = uv2mesh(static[np.newaxis,:,:,:]).numpy()[0]
+    mesh_body = uv2mesh(body[np.newaxis,:,:,:]).numpy()[0]
+    fig = make_subplots(rows=1, cols=2, start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(2)] ])
     F = np.load("custom_utils/smpl_faces.npy")
     _, faces = _maskFaces(np.ones((14475, 3)), F, mesh_mask)
-    pass
+
+    x = mesh_static[mesh_mask, :]
+    y = mesh_body
+    fig.add_trace(go.Mesh3d(
+                    x=x[:,0],
+                    y=x[:,1],
+                    z=x[:,2],
+                    i=faces[:,0],
+                    j=faces[:,1],
+                    k=faces[:,2]
+                ),
+                row=1, col=1)
+    fig.add_trace(go.Mesh3d(
+                    x=y[:,0],
+                    y=y[:,2],
+                    z=y[:,1],
+                    i=F[:,0],
+                    j=F[:,1],
+                    k=F[:,2]
+                ),
+                row=1, col=2)
+    make_square_2(fig)
+    fig.show()
+
+
+def plot_variable_condition(condition, mesh_mask, is_curv=False, is_body=False):
+    if is_curv:
+        img = np.concatenate(condition.argmax(-1), axis=-1)
+        ex.imshow(img / img.max()).show()
+    else:
+        mesh = uv2mesh(condition).numpy()
+        fig = make_subplots(rows=1, cols=condition.shape[0], start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(condition.shape[0])] ])
+        F = np.load("custom_utils/smpl_faces.npy")
+        if is_body:
+            faces = [F for _ in range(condition.shape[0])]
+        else:
+            _, faces = maskFacesBatch(np.ones((condition.shape[0], 14475, 3)), F, mesh_mask)
+        for i, m in enumerate(mesh):
+            if not is_body:
+                m = m[mesh_mask[i], :]
+            if is_body:
+                x = m.copy()
+                m[:,1] = x[:,2]
+                m[:,2] = x[:,1]
+            fig.add_trace(go.Mesh3d(
+                            # 8 vertices of a cube
+                            x=m[:,0],
+                            y=m[:,1],
+                            z=m[:,2],
+                            # i, j and k give the vertices of triangles
+                            i = faces[i][:,0],
+                            j = faces[i][:,1],
+                            k = faces[i][:,2]
+                        ),
+                        row=1, col=i+1)
+        make_square_4(fig)
+        fig.show()
+
+
+def plot_results(unposed, posed, body, mesh_mask):
+    fig = make_subplots(rows=1, cols=unposed.shape[0], start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(unposed.shape[0])] ])
+    F = np.load("custom_utils/smpl_faces.npy")
+    _, faces = maskFacesBatch(np.ones((unposed.shape[0], 14475, 3)), F, mesh_mask)
+    body_mesh = uv2mesh(body).numpy()
+
+    collision = [np.concatenate((body_mesh[i], posed[i][mesh_mask[i], :]), axis=0) for i in range(unposed.shape[0])]
+    collision_faces = [np.concatenate((F, faces[i] + 14475), axis=0) for i in range(unposed.shape[0])]
+    collision_color = [np.concatenate((np.ones_like(body_mesh[i], dtype=np.int32)*100, np.ones_like(unposed[i], dtype=np.int32)*230), axis=0) for i in range(unposed.shape[0])]
+
+    for i, c in enumerate(collision):
+        fig.add_trace(go.Mesh3d(
+                        x=c[:,0],
+                        y=c[:,1],
+                        z=c[:,2],
+                        i=collision_faces[i][:,0],
+                        j=collision_faces[i][:,1],
+                        k=collision_faces[i][:,2],
+                        vertexcolor=collision_color[i]
+                    ),
+                    row=1, col=i+1)
+    make_square_4(fig)
+    fig.show()
+
+
+
+def get_noise(fix_noise, nr, latent_dim=128):
+    eps = tf.random.normal(shape=(1 if fix_noise else nr, latent_dim))
+    if fix_noise:
+        eps = tf.tile(eps, [nr,1])
+    return eps
+
 
 
 def sample(folder, variable_condition, fix_noise, nr=4):
     # run = wandb.init(project='vae_plots', name=)
-
     random_static, random_body, random_curv, random_unposed = variable_condition == 'static', variable_condition == 'body', variable_condition == 'curv', variable_condition == 'unposed'
 
-    frame_main, frame_rand = get_frame(folder, nr)
+    frames = get_frame(folder, nr)
-    unposed, uv_mask, mesh_mask = get_unposed(folder, frame, nr, random_unposed)
+    noise = get_noise(fix_noise, nr)
-    static = get_statics(folder, nr, random_static)
+    noise_curv = get_noise(fix_noise, nr)
-    body, pose, shape = get_bodies(folder, nr, random_body)
+    body, pose_shape = get_bodies(folder, frames, random_body)
-    curv = get_curvs(folder, body, static, uv_mask, frame, nr, random_curv, curv_model)
+    static, uv_mask, mesh_mask = get_statics(folder, frames, random_static)
+    unposed, mesh_unposed = get_unposed(folder, frames, random_unposed)
+    curv = get_curvs(folder, body, static, uv_mask, frames, noise_curv, random_curv, curv_model, use_original=True)
 
-    pred = model([unposed, static, body, curv, uv_mask], fix_noise=fix_noise)
+    pred = model.sample(static, body, curv, uv_mask, noise=noise)
     mesh_pred = uv2mesh(pred)
 
-    gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
+    gender = tf.cast(tf.greater_equal(pose_shape[:, -1] , 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])])
+    shape_offsets = tf.stack([tf.gather(tf.tensordot(pose_shape[i,72:-1][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(len(frames))])
+    mesh_pred_shaped = (mesh_pred + shape_offsets).numpy()
+    mesh_posed = smpl.run_garment(pose_shape[:, :72], pose_shape[:, 72:], mesh_pred_shaped).numpy()
 
-    mesh_pred_shaped = mesh_pred + shape_offsets
+    plot_inputs(unposed[0], static[0], body[0], curv[0], mesh_mask[0])
 
-    mesh_posed = smpl.run_garment(pose, shape, mesh_pred_shaped).numpy()
+    variable_cond_data = None
-
-    plot_inputs(unposed[0], static[0], body[0], curv[0])
-
-    variable_cond = None
     if random_static:
-        variable_cond = static
+        variable_cond_data = static
     elif random_body:
-        variable_cond = body
+        variable_cond_data = body
     elif random_curv:
-        variable_cond = curv
+        variable_cond_data = curv
     elif random_unposed:
-        variable_cond = unposed
+        variable_cond_data = unposed
     else:
         assert False
 
-    plot_variable_condition(variable_cond, is_curv=random_curv)
+    plot_variable_condition(variable_cond_data, mesh_mask, is_curv=random_curv, is_body=random_body)
 
-    plot_results(mesh_pred, mesh_posed, body, mesh_mask)
+    plot_results(mesh_pred_shaped, mesh_posed, body, mesh_mask)
 
 
     print()
@@ -256,8 +407,8 @@ if __name__ == '__main__':
 
     folder = '/home/workspace/local-data/train_100'
     folder_main = 'checkpoints/cvae_13'
-    folder_curv = 'checkpoints/cvae_curv_2'
+    folder_curv = 'checkpoints/cvae_curv_3'
-    variable_condition = 'curv'
+    variable_condition = 'static'
     batch_size = 4
     
     smpl = SMPL()