large refactor

sample_pro.py

@@ -1,5 +1,5 @@
 import os
-os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+os.environ["CUDA_VISIBLE_DEVICES"] = "3"
 import tensorflow as tf
 
 gpus = tf.config.experimental.list_physical_devices('GPU')
@@ -27,6 +27,7 @@ from PIL import Image
 from plotly.subplots import make_subplots
 
 from custom_utils.smpl import SMPL
+from custom_utils.custom_utils import load_model
 from cvae_curv import CVAE as CVAECurv
 from models.cvae_model import CVAE as CVAEMain
 
@@ -136,12 +137,6 @@ 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 get_frame(folder, nr):
     pose_folder = os.path.join(folder, 'pose_shape') 
     frames = sorted(os.listdir(pose_folder))
@@ -391,6 +386,7 @@ def sample(folder, variable_condition, fix_noise, nr=4):
         assert False
 
     plot_variable_condition(variable_cond_data, mesh_mask, is_curv=random_curv, is_body=random_body)
+    plot_variable_condition(curv, mesh_mask, is_curv=True, is_body=False)
 
     plot_results(mesh_pred_shaped, mesh_posed, body, mesh_mask)
 
@@ -407,7 +403,7 @@ if __name__ == '__main__':
     args = vars(parser.parse_args())
 
     folder = '/home/workspace/local-data/train_100'
-    folder_main = 'checkpoints/cvae_13'
+    folder_main = 'checkpoints/cvae_18_copy'
     folder_curv = 'checkpoints/cvae_curv_3'
     variable_condition = 'static'
     batch_size = 4
@@ -415,13 +411,17 @@ if __name__ == '__main__':
     smpl = SMPL()
 
     model = CVAEMain(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_8',
-                    curvature_model='checpoints/ae_curvature_7')
-    curv_model = CVAECurv()
+                    latent_dim=128, static_model='checkpoints/ae_static_6_copy',
+                    body_model='checkpoints/ae_body_8_copy',
+                    curvature_model='checkpoints/ae_curvature_7_copy',
+                    freeze_encoders=True)
+    curv_model = CVAECurv(static_model='checkpoints/ae_static_6_copy',
+                          body_model='checkpoints/ae_body_8_copy',)
 
 
     load_model(model, folder_main)
     load_model(curv_model, folder_curv)
 
     sample(folder, variable_condition, fix_noise=True)
+
+    print()