pyramid with common encoders per level

cvae_pyramid.py

@@ -1,5 +1,5 @@
 import os
-os.environ["CUDA_VISIBLE_DEVICES"] = "2"
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
 import tensorflow as tf
 
 gpus = tf.config.experimental.list_physical_devices('GPU')
@@ -37,7 +37,7 @@ from custom_utils.smpl import SMPL
 from models.cvae_pyramid_model import Pyramid
 
 
-name='cvae_pyr_1_d128'
+name='cvae_pyr_2_d128'
 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. \
@@ -48,7 +48,8 @@ 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\
 training on 100 garments with offsets over the average of these\
-l2 loss on both 128 and full res, bottom predicts base, high res predicts offsets over this')
+l2 loss on both 128 and full res, bottom predicts base, high res predicts offsets over this\
+low and high res both use same full res condition encoders, but low res also gets low res versions for main encoder input')
     artifact = wandb.Artifact('code', type='code')
     artifact.add_file('cvae_pyramid.py')
     artifact.add_file('models/cvae_model.py')
@@ -250,19 +251,21 @@ def plot_everything(batch: dict, type: str):
     img = tf.concat((uv_static_img, uv_body_posed_img, orig_img, pred_img, pred_low_up), axis=2)
     imgs = [wandb.Image((img[i] + 1) / 2.0) for i in range(4)]
 
+    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])])
+
     F = np.load("custom_utils/smpl_faces.npy")
     mask = batch['mesh_mask'].numpy().astype(bool)
     garment_unposed, faces = maskFacesBatch(mesh_pred_full.numpy(), F, mask)
     garment_unposed_low, faces = maskFacesBatch(mesh_pred_low.numpy(), F, mask)
+    garment_unposed = [garment_unposed[i] + shape_offsets.numpy()[i, mask[i], :] for i in range(mask.shape[0])]
+    garment_unposed_low = [garment_unposed_low[i] + shape_offsets.numpy()[i, mask[i], :] for i in range(mask.shape[0])]
     gt = [mesh_orig[i].numpy()[mask[i], :] for i in range(pred_full.shape[0])]
     
     body = batch['mesh_body_posed'].numpy()
     mesh_posed_full = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_full).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_full[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))]
@@ -439,7 +442,12 @@ if __name__ == '__main__':
     static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/valid', signature=signature_static)
 
 
-    model = Pyramid(latent_dim=128, low_res=low_resolution)
+    model = Pyramid(latent_dim=128,
+                    low_res=low_resolution, 
+                    freeze_encoders=True, 
+                    static_model='checkpoints/ae_static_4',
+                    body_model='checpoints/ae_body_7',
+                    curvature_model='checpoints/ae_curvature_7')
 
     # 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)

models/cvae_model.py

@@ -258,21 +258,31 @@ class CVAE32(tf.keras.Model):
                  static_model='',
                  body_model='',
                  curvature_model='', 
-                 freeze_encoders=False):
+                 freeze_encoders=False,
+                 use_full_res_encoders=True):
         super(CVAE32, self).__init__()
         self.latent_dim = latent_dim
 
-        model : StaticAE32 = StaticAE32()
+        if use_full_res_encoders:
+            model : StaticAE = StaticAE()
+        else:
+            model : StaticAE32 = StaticAE32()
         if static_model != '':
             load_model(model, static_model)
         self.static_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
 
-        model : BodyAE32 = BodyAE32()
+        if use_full_res_encoders:
+            model : BodyAE = BodyAE()
+        else:
+            model : BodyAE32 = BodyAE32()
         if body_model != '':
             load_model(model, body_model)
         self.body_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
 
-        model : CurvatureAE32 = CurvatureAE32()
+        if use_full_res_encoders:
+            model : CurvatureAE = CurvatureAE()
+        else:
+            model : CurvatureAE32 = CurvatureAE32()
         if curvature_model != '':
             load_model(model, curvature_model)
         self.curv_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
@@ -429,21 +439,31 @@ class CVAE128(tf.keras.Model):
                  static_model='',
                  body_model='',
                  curvature_model='', 
-                 freeze_encoders=False):
+                 freeze_encoders=False,
+                 use_full_res_encoders=True):
         super(CVAE128, self).__init__()
         self.latent_dim = latent_dim
 
-        model : StaticAE128 = StaticAE128()
+        if use_full_res_encoders:
+            model : StaticAE = StaticAE()
+        else:
+            model : StaticAE128 = StaticAE128()
         if static_model != '':
             load_model(model, static_model)
         self.static_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
-
+        
-        model : BodyAE128 = BodyAE128()
+        if use_full_res_encoders:
+            model : BodyAE = BodyAE()
+        else:
+            model : BodyAE128 = BodyAE128()
         if body_model != '':
             load_model(model, body_model)
         self.body_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
 
-        model : CurvatureAE128 = CurvatureAE128()
+        if use_full_res_encoders:
+            model : CurvatureAE = CurvatureAE()
+        else:
+            model : CurvatureAE128 = CurvatureAE128()
         if curvature_model != '':
             load_model(model, curvature_model)
         self.curv_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
@@ -515,11 +535,11 @@ class CVAE128(tf.keras.Model):
 
 
     def call(self, inputs, training=None, zero_latent=False):
-        uv_unposed, uv_template, uv_body_posed, uv_curv = inputs
+        uv_unposed, uv_template, uv_body_posed, uv_curv, uv_template_high, uv_body_posed_high, uv_curv_high = 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_static = self.static_enc(uv_template_high)
-        z_body = self.body_enc(uv_body_posed)
+        z_body = self.body_enc(uv_body_posed_high)
-        z_curv = self.curv_enc(uv_curv)
+        z_curv = self.curv_enc(uv_curv_high)
 
         x = tf.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], axis=-1)
 

models/cvae_pyramid_model.py

@@ -12,15 +12,39 @@ from models.cvae_model import CVAE, CVAE128, CVAE32
 
 class Pyramid(tf.keras.Model):
     
-    def __init__(self, low_res=32, latent_dim=128, avg_garment_path=''):
+    def __init__(self, low_res=32,
+                 latent_dim=128,
+                 avg_garment_path='',
+                 static_model='',
+                 body_model='',
+                 curvature_model='', 
+                 freeze_encoders=False,
+                 use_full_res_encoders=True):
         super(Pyramid, self).__init__()
 
         if low_res == 32:
-            self.model_low = CVAE32(latent_dim=latent_dim)
+            self.model_low = CVAE32(latent_dim=latent_dim,
+                                    static_model=static_model, 
+                                    body_model=body_model, 
+                                    curvature_model=curvature_model, 
+                                    freeze_encoders=freeze_encoders, 
+                                    use_full_res_encoders=use_full_res_encoders)
         if low_res == 128:
-            self.model_low = CVAE128(latent_dim=latent_dim)
+            self.model_low = CVAE128(latent_dim=latent_dim,
-        self.model_full = CVAE(latent_dim=latent_dim)
+                                    static_model=static_model, 
+                                    body_model=body_model, 
+                                    curvature_model=curvature_model, 
+                                    freeze_encoders=freeze_encoders, 
+                                    use_full_res_encoders=use_full_res_encoders)
+
+        self.model_full = CVAE(latent_dim=latent_dim,
+                               static_model=static_model, 
+                               body_model=body_model, 
+                               curvature_model=curvature_model, 
+                               freeze_encoders=freeze_encoders)
         
+        self.use_full_res_encoders = use_full_res_encoders
+
         if avg_garment_path != '':
             self.avg_garment = np.load(avg_garment_path)
         else:
@@ -29,7 +53,10 @@ class Pyramid(tf.keras.Model):
     def call(self, inputs, training=None):
         uv_unposed, uv_static, uv_body_posed, uv_curv, uv_unposed_low, uv_static_low, uv_body_posed_low, uv_curv_low = inputs
 
-        pred_low, mean_low, logvar_low, _ = self.model_low([uv_unposed_low, uv_static_low, uv_body_posed_low, uv_curv_low], training=training)
+        if self.use_full_res_encoders:
+            pred_low, mean_low, logvar_low, _ = self.model_low([uv_unposed_low, uv_static_low, uv_body_posed_low, uv_curv_low, uv_static, uv_body_posed, uv_curv], training=training)
+        else:
+            pred_low, mean_low, logvar_low, _ = self.model_low([uv_unposed_low, uv_static_low, uv_body_posed_low, uv_curv_low], training=training)
 
         x = tf.image.resize(pred_low, (512, 256), tf.image.ResizeMethod.LANCZOS5)