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)