commit all

.gitignore

@@ -146,3 +146,5 @@ wandb copy/
 checkpoints/
 temp/
 ids/
+examples/
+custom_data/dijkstra*

analysis_curvature.py

@@ -0,0 +1,114 @@
+import os
+import sys
+from configs.config_abl_skin_var_3 import config
+
+os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
+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 glob
+import time
+import argparse
+import wandb
+import tqdm
+import random
+import numpy as np
+import plotly.graph_objects as go
+import plotly.express as ex
+from PIL import Image
+from postprocessing import fix_collision
+
+if config['is_skirt']:
+    from losses import edge_loss_skirt as edge_loss
+else:
+    from losses import edge_loss
+
+
+
+from losses import collision_ratio, kl_loss #, concept_loss
+# from losses import collision_penalty as collision_penalty
+from losses import collision_loss as collision_penalty
+from custom_utils.custom_utils import batch_to_dict, load_model, plot_3d, UV2Mesh, get_summary_html, maskFacesBatch, _maskFaces, plot_3d_collision, UV2MeshLabel
+from dataset import DatasetFactory, StaticDataset
+from custom_utils.smpl import SMPL
+from models.cvae_model import CVAE
+from custom_utils.supersmpl_io import writeOBJ
+
+
+uv2mesh = UV2MeshLabel(res=256)
+
+
+@tf.function
+def step(batch):
+    pred, mean, logvar = model([batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']])
+
+    mesh_pred = tf.math.argmax(uv2mesh(pred), axis=-1)
+
+    return mesh_pred
+
+
+if __name__ == '__main__':
+    # init_wandb()
+    parser = argparse.ArgumentParser(description='')
+    parser.add_argument('--name', type=str)
+    args = vars(parser.parse_args())
+
+    smpl = SMPL(use_one_hot=config['use_one_hot'])
+    skirt_modifier = '_skirt' if config['is_skirt'] else ''
+
+    folder = config['test_dir']
+    smpl = SMPL(use_one_hot=False)
+    model = CVAE(config)
+
+    # str_split = config['name'].split('_')
+    # test_idx = str_split.index('test')
+    # checkpoint_name = '_'.join(str_split[:test_idx])
+    checkpoint_name = config['name']
+
+    F = np.load("custom_utils/smpl_faces" + skirt_modifier + ".npy").astype(np.int32)
+
+    ids = os.listdir(os.path.join(folder, 'pose_shape'))
+    ids = [id.split('_')[0] for id in ids]
+    ids = set(ids)
+
+    iteration = '70000'
+
+    load_model(model, 'checkpoints/' + config['name'], iteration=iteration)
+
+
+    factory = DatasetFactory()
+    dataset_test = factory.get_dataset(folder=config['test_dir'], 
+                                       batch_size=1, 
+                                       signature=config['signature_dynamic'], 
+                                       resolutions=[config['resolution']], 
+                                       in_memory=False, 
+                                       do_repeat=False,
+                                       use_one_hot=config['use_one_hot'],
+                                       use_no_orient=config['remove_global_orientation'],
+                                       frames_file='/home/workspace/data/dataset/test_ready/frames_pose_cluster_20_sample_10.txt')
+    static_dataset_test = StaticDataset(folder=config['test_dir'], 
+                                        signature=config['signature_static'], 
+                                        resolutions=[config['resolution']], 
+                                        in_memory=True)
+
+    # faces_folder = '/home/workspace/data/results/support/faces' + ('_skirt' if config['is_skirt'] else '')
+    # ids = os.listdir(faces_folder)
+    # ids = [id.split('.')[0] for id in ids]
+    # faces = {id : np.load(faces_folder + '/' + id + '.npy') for id in ids }
+
+    for batch in tqdm.tqdm(dataset_test):
+        batch = static_dataset_test.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
+
+
+

auto_encoder.py

@@ -124,7 +124,7 @@ class Bottleneck(tf.keras.Model):
                        layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
             )
 
-    def call(self, inputs, training=None):
+    def call(self, inputs, training):
         x = inputs
 
         x = self.bottleneck(x, training=training)
@@ -145,7 +145,7 @@ class AutoEncoder(tf.keras.Model):
         self.decoder = ConvNeXtReverse(depths=list(reversed([3, 3, 9, 3])), dims=list(reversed([bottleneck_size, 128, 64, 16])))
 
 
-    def call(self, inputs, training=None):
+    def call(self, inputs, training):
         x, uv_mask = inputs
 
         x = tf.transpose(x, [0,2,3,1])
@@ -328,7 +328,7 @@ if __name__ == '__main__':
 
 
     factory = DatasetFactory()
-    signature_dynamic = { 'id': (),
+    signature_dynamic = { 'id': (), 'frame_nr': (),
                           'pose_shape': (72 + 11,),
                         #   'uv_posed': (3, 512, 256,),
                           'uv_body_posed': (3, 512, 256,),

auto_encoder_body.py

@@ -1,5 +1,6 @@
 import os
-from configs.config_ae_body import config
+import sys
+from configs.config_ae_body_post import config
 os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
 
 import tensorflow as tf
@@ -29,7 +30,7 @@ from tf_utils import check_tf_version, allow_memory_growth
 from losses import edge_loss, uv_mask_loss
 from dataset import DatasetFactory, StaticDataset
 from custom_utils.custom_utils import batch_to_dict, plot_3d
-from models.cond_encoder import AutoEncoder
+from models.cond_encoder import AutoEncoderNew as AutoEncoder
 
 
 def init_wandb():
@@ -37,12 +38,13 @@ def init_wandb():
     artifact = wandb.Artifact('code', type='code')
     artifact.add_file('auto_encoder_body.py')
     artifact.add_file('losses.py')
-    artifact.add_file('configs/config_base.py')
-    artifact.add_file('configs/config_ae_body.py')
     artifact.add_file('models/base.py')
     artifact.add_file('models/cond_encoder.py')
     artifact.add_file('dataset.py')
     artifact.add_file('convnext/models/convnext_tf.py')
+    config_file_names = [key for key in sys.modules.keys()  if key.startswith('configs.')]
+    for config_file_name in config_file_names:
+        artifact.add_file('/'.join(config_file_name.split('.')) + '.py')
     run.log_artifact(artifact)
 
 uv2mesh = UV2Mesh(res=256, is_smpl_orig=config['is_cape'], is_skirt=config['is_skirt'])
@@ -164,10 +166,33 @@ if __name__ == '__main__':
 
     factory = DatasetFactory()
 
-    dataset_train = factory.get_dataset(folder=config['train_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=config['dataset_in_memory'], is_cape=config['is_cape'])
+    dataset_train = factory.get_dataset(folder=config['train_dir'],
-    static_dataset_train = StaticDataset(folder=config['train_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=True, is_cape=config['is_cape'])
+                                        batch_size=config['batch_size'], 
-    dataset_valid = factory.get_dataset(folder=config['valid_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=config['dataset_in_memory'], is_cape=config['is_cape'])
+                                        signature=config['signature_dynamic'],
-    static_dataset_valid = StaticDataset(folder=config['valid_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=True, is_cape=config['is_cape'])
+                                        resolutions=[config['resolution']], 
+                                        in_memory=config['dataset_in_memory'],
+                                        use_no_orient=config['remove_global_orientation'],
+                                        frames_file=config['frames_file'] if 'frames_file' in config.keys() else None, 
+                                        is_cape=config['is_cape'])
+    static_dataset_train = StaticDataset(folder=config['train_dir'],
+                                         signature=config['signature_static'], 
+                                         resolutions=[config['resolution']], 
+                                         in_memory=True, 
+                                        is_cape=config['is_cape'])
+
+    dataset_valid = factory.get_dataset(folder=config['valid_dir'],
+                                        batch_size=config['batch_size'], 
+                                        signature=config['signature_dynamic'], 
+                                        resolutions=[config['resolution']], 
+                                        in_memory=config['dataset_in_memory'], 
+                                        use_no_orient=config['remove_global_orientation'], 
+                                        frames_file=config['frames_file_valid'] if 'frames_file_valid' in config.keys() else None, 
+                                        is_cape=config['is_cape'])
+    static_dataset_valid = StaticDataset(folder=config['valid_dir'],
+                                         signature=config['signature_static'], 
+                                         resolutions=[config['resolution']], 
+                                         in_memory=True, 
+                                        is_cape=config['is_cape'])
 
     model = AutoEncoder(config['cond_ae'])
     optimizer = tf.optimizers.Adam(learning_rate=config['lr'])
@@ -176,8 +201,8 @@ if __name__ == '__main__':
     # load_model(model.decoder, 'checkpoints/body_1_skirt_2_dec')
 
     ckpt_enc = tf.train.Checkpoint(model=model.encoder)
-    manager_enc = tf.train.CheckpointManager(ckpt_enc, 'checkpoints/' + config['name'] + '_enc', max_to_keep=20)
+    manager_enc = tf.train.CheckpointManager(ckpt_enc, 'checkpoints/post/' + config['name'] + '_enc', max_to_keep=20)
     ckpt_dec = tf.train.Checkpoint(model=model.decoder)
-    manager_dec = tf.train.CheckpointManager(ckpt_dec, 'checkpoints/' + config['name'] + '_dec', max_to_keep=20)
+    manager_dec = tf.train.CheckpointManager(ckpt_dec, 'checkpoints/post/' + config['name'] + '_dec', max_to_keep=20)
 
     train()

auto_encoder_curv.py

@@ -1,5 +1,6 @@
 import os
-from configs.config_ae_curv import config
+import sys
+from configs.config_ae_curv_post import config
 os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
 
 import tensorflow as tf
@@ -20,6 +21,7 @@ import argparse
 import wandb
 import random
 import numpy as np
+import cmapy as cm
 
 
 from custom_utils.custom_utils import UV2Mesh, UV2MeshLabel, get_summary_html, maskFacesBatch
@@ -29,7 +31,7 @@ import plotly.express as ex
 from tf_utils import check_tf_version, allow_memory_growth
 from dataset import DatasetFactory, StaticDataset
 from custom_utils.custom_utils import batch_to_dict
-from models.cond_encoder import AutoEncoder
+from models.cond_encoder import AutoEncoderNew as AutoEncoder
 
 
 testing = False
@@ -39,12 +41,13 @@ def init_wandb():
         artifact = wandb.Artifact('code', type='code')
         artifact.add_file('auto_encoder_curv.py')
         artifact.add_file('losses.py')
-        artifact.add_file('configs/config_base.py')
-        artifact.add_file('configs/config_ae_curv.py')
         artifact.add_file('models/base.py')
         artifact.add_file('models/cond_encoder.py')
         artifact.add_file('dataset.py')
         artifact.add_file('convnext/models/convnext_tf.py')
+        config_file_names = [key for key in sys.modules.keys()  if key.startswith('configs.')]
+        for config_file_name in config_file_names:
+            artifact.add_file('/'.join(config_file_name.split('.')) + '.py')
         run.log_artifact(artifact)
 
 uv2mesh = UV2MeshLabel(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
@@ -99,16 +102,20 @@ def train_step(batch):
 
         ce_uv_loss = scce(tf.math.argmax(batch['uv_curvature'], axis=-1), pred)
         ms_ssim = tf.reduce_mean(tf.image.ssim(pred, batch['uv_curvature'], max_val=1))
-        loss = ce_uv_loss
 
-    pred_mesh = tf.math.argmax(uv2mesh(pred), axis=-1)
+        mesh_pred = uv2mesh(pred) * batch['mesh_mask'][:,:,tf.newaxis]
+        mesh_orig = tf.cast(uv2mesh(tf.math.argmax(batch['uv_curvature'], axis=-1, output_type=tf.int32)), tf.float32) * batch['mesh_mask']
+        ce_mesh_loss = scce(mesh_orig, mesh_pred)
+
+        loss = ce_uv_loss + ce_mesh_loss
+
     mean_iou = get_mean_iou(tf.math.argmax(batch['uv_curvature'], axis=-1), tf.math.argmax(pred, axis=-1))
-    mean_iou_mesh = get_mean_iou(tf.cast(batch['mesh_curvature'], tf.float32) * batch['mesh_mask'], tf.cast(pred_mesh, tf.float32) * batch['mesh_mask'])
     dice = dice_dense(pred, batch['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/mean_iou_mesh': mean_iou_mesh, 'train/dice': dice }
+    return { 'train/loss': loss, 'train/ce_uv_loss': ce_uv_loss, 'train/ssim': ms_ssim, 'train/mean_iou': mean_iou,
+             'train/dice': dice, 'train/ce_mesh_loss': ce_mesh_loss }
         
 
 
@@ -118,14 +125,17 @@ def valid_step(batch):
 
     ce_uv_loss = scce(tf.math.argmax(batch['uv_curvature'], axis=-1), pred)
     ms_ssim = tf.reduce_mean(tf.image.ssim(pred, batch['uv_curvature'], max_val=1))
-    loss = ce_uv_loss
 
-    pred_mesh = tf.math.argmax(uv2mesh(pred), axis=-1)
+    mesh_pred = uv2mesh(pred) * batch['mesh_mask'][:,:,tf.newaxis]
+    mesh_orig = tf.cast(uv2mesh(tf.math.argmax(batch['uv_curvature'], axis=-1, output_type=tf.int32)), tf.float32) * batch['mesh_mask']
+    ce_mesh_loss = scce(mesh_orig, mesh_pred)
+    loss = ce_uv_loss + ce_mesh_loss
+
     mean_iou = get_mean_iou(tf.math.argmax(batch['uv_curvature'], axis=-1), tf.math.argmax(pred, axis=-1))
-    mean_iou_mesh = get_mean_iou(tf.cast(batch['mesh_curvature'], tf.float32) * batch['mesh_mask'], tf.cast(pred_mesh, tf.float32) * batch['mesh_mask'])
     dice = dice_dense(pred, batch['uv_curvature'])
 
-    return { 'valid/loss': loss, 'valid/ce_uv_loss': ce_uv_loss, 'valid/ssim': ms_ssim, 'valid/mean_iou': mean_iou, 'valid/mean_iou_mesh': mean_iou_mesh, 'valid/dice': dice }
+    return { 'valid/loss': loss, 'valid/ce_uv_loss': ce_uv_loss, 'valid/ssim': ms_ssim, 'valid/mean_iou': mean_iou,
+             'valid/dice': dice, 'valid/ce_mesh_loss': ce_mesh_loss }
         
 
 def plot_everything(batch: dict, type: str):
@@ -134,7 +144,14 @@ def plot_everything(batch: dict, type: str):
     orig_img = tf.math.argmax(batch['uv_curvature'], axis=-1, output_type=tf.dtypes.int32)
     pred_img = tf.math.argmax(pred, axis=-1, output_type=tf.dtypes.int32)
     img = tf.concat((orig_img, pred_img), axis=2).numpy()
-    imgs = { type + '/images': [wandb.Image(img[i,:,:,np.newaxis] / 6.0) for i in range(img.shape[0])] }
+
+    colors = np.array([cm.color('jet', 1 - i) for i in np.linspace(0, 1, 6+1)])
+    color_img = np.zeros((img.shape[0], img.shape[1], img.shape[2], 3), np.uint8)
+    for i in range(7):
+        color_img[img == i, :] = colors[i]
+    color_img = [wandb.Image(color_img[i]) for i in range(4)]
+
+    imgs = { type + '/images': [wandb.Image(img[i,:,:,np.newaxis] / 6.0) for i in range(img.shape[0])], type + '/color_images': color_img }
 
     return imgs
 
@@ -202,18 +219,41 @@ if __name__ == '__main__':
 
     factory = DatasetFactory()
 
-    dataset_train = factory.get_dataset(folder=config['train_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=config['dataset_in_memory'], is_cape=config['is_cape'])
+    dataset_train = factory.get_dataset(folder=config['train_dir'],
-    static_dataset_train = StaticDataset(folder=config['train_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=True, is_cape=config['is_cape'])
+                                        batch_size=config['batch_size'], 
-    dataset_valid = factory.get_dataset(folder=config['valid_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=config['dataset_in_memory'], is_cape=config['is_cape'])
+                                        signature=config['signature_dynamic'],
-    static_dataset_valid = StaticDataset(folder=config['valid_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=True, is_cape=config['is_cape'])
+                                        resolutions=[config['resolution']], 
+                                        in_memory=config['dataset_in_memory'],
+                                        use_no_orient=config['remove_global_orientation'],
+                                        frames_file=config['frames_file'] if 'frames_file' in config.keys() else None, 
+                                        is_cape=config['is_cape'])
+    static_dataset_train = StaticDataset(folder=config['train_dir'],
+                                         signature=config['signature_static'], 
+                                         resolutions=[config['resolution']], 
+                                         in_memory=True, 
+                                        is_cape=config['is_cape'])
+
+    dataset_valid = factory.get_dataset(folder=config['valid_dir'],
+                                        batch_size=config['batch_size'], 
+                                        signature=config['signature_dynamic'], 
+                                        resolutions=[config['resolution']], 
+                                        in_memory=config['dataset_in_memory'], 
+                                        use_no_orient=config['remove_global_orientation'], 
+                                        frames_file=config['frames_file_valid'] if 'frames_file_valid' in config.keys() else None, 
+                                        is_cape=config['is_cape'])
+    static_dataset_valid = StaticDataset(folder=config['valid_dir'],
+                                         signature=config['signature_static'], 
+                                         resolutions=[config['resolution']], 
+                                         in_memory=True, 
+                                        is_cape=config['is_cape'])
 
     model = AutoEncoder(config['cond_ae'])
     optimizer = tf.optimizers.Adam(learning_rate=config['lr'])
 
     ckpt_enc = tf.train.Checkpoint(model=model.encoder)
-    manager_enc = tf.train.CheckpointManager(ckpt_enc, 'checkpoints/' + config['name'] + '_enc', max_to_keep=20)
+    manager_enc = tf.train.CheckpointManager(ckpt_enc, 'checkpoints/post/' + config['name'] + '_enc', max_to_keep=20)
     ckpt_dec = tf.train.Checkpoint(model=model.decoder)
-    manager_dec = tf.train.CheckpointManager(ckpt_dec, 'checkpoints/' + config['name'] + '_dec', max_to_keep=20)
+    manager_dec = tf.train.CheckpointManager(ckpt_dec, 'checkpoints/post/' + config['name'] + '_dec', max_to_keep=20)
 
     if testing:
         test()

auto_encoder_static.py

@@ -1,4 +1,5 @@
 import os
+import sys
 from configs.config_ae_static import config
 os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
 
@@ -32,7 +33,7 @@ else:
 from tf_utils import check_tf_version, allow_memory_growth
 from dataset import DatasetFactory, StaticDataset
 from custom_utils.custom_utils import batch_to_dict, plot_3d
-from models.cond_encoder import AutoEncoder
+from models.cond_encoder import AutoEncoderNew as AutoEncoder
 
 
 def init_wandb():
@@ -40,12 +41,13 @@ def init_wandb():
     artifact = wandb.Artifact('code', type='code')
     artifact.add_file('auto_encoder_static.py')
     artifact.add_file('losses.py')
-    artifact.add_file('configs/config_base.py')
-    artifact.add_file('configs/config_ae_static.py')
     artifact.add_file('models/base.py')
     artifact.add_file('models/cond_encoder.py')
     artifact.add_file('dataset.py')
     artifact.add_file('convnext/models/convnext_tf.py')
+    config_file_names = [key for key in sys.modules.keys()  if key.startswith('configs.')]
+    for config_file_name in config_file_names:
+        artifact.add_file('/'.join(config_file_name.split('.')) + '.py')
     run.log_artifact(artifact)
 
 uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
@@ -184,7 +186,7 @@ def test():
 
 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())
@@ -198,10 +200,33 @@ if __name__ == '__main__':
 
     factory = DatasetFactory()
 
-    dataset_train = factory.get_dataset(folder=config['train_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=config['dataset_in_memory'], is_cape=config['is_cape'])
+    dataset_train = factory.get_dataset(folder=config['train_dir'],
-    static_dataset_train = StaticDataset(folder=config['train_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=True, is_cape=config['is_cape'])
+                                        batch_size=config['batch_size'], 
-    dataset_valid = factory.get_dataset(folder=config['valid_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=config['dataset_in_memory'], is_cape=config['is_cape'])
+                                        signature=config['signature_dynamic'],
-    static_dataset_valid = StaticDataset(folder=config['valid_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=True, is_cape=config['is_cape'])
+                                        resolutions=[config['resolution']], 
+                                        in_memory=config['dataset_in_memory'],
+                                        use_no_orient=config['remove_global_orientation'],
+                                        frames_file=config['frames_file'] if 'frames_file' in config.keys() else None, 
+                                        is_cape=config['is_cape'])
+    static_dataset_train = StaticDataset(folder=config['train_dir'],
+                                         signature=config['signature_static'], 
+                                         resolutions=[config['resolution']], 
+                                         in_memory=True, 
+                                        is_cape=config['is_cape'])
+
+    dataset_valid = factory.get_dataset(folder=config['valid_dir'],
+                                        batch_size=config['batch_size'], 
+                                        signature=config['signature_dynamic'], 
+                                        resolutions=[config['resolution']], 
+                                        in_memory=config['dataset_in_memory'], 
+                                        use_no_orient=config['remove_global_orientation'], 
+                                        frames_file=config['frames_file_valid'] if 'frames_file_valid' in config.keys() else None, 
+                                        is_cape=config['is_cape'])
+    static_dataset_valid = StaticDataset(folder=config['valid_dir'],
+                                         signature=config['signature_static'], 
+                                         resolutions=[config['resolution']], 
+                                         in_memory=True, 
+                                        is_cape=config['is_cape'])
 
     model : AutoEncoder = AutoEncoder(config['cond_ae'])
     optimizer = tf.optimizers.Adam(learning_rate=config['lr'])
@@ -211,9 +236,9 @@ if __name__ == '__main__':
         load_model(model.decoder, config['pretrained_model'] + '_dec')
 
     ckpt_enc = tf.train.Checkpoint(model=model.encoder)
-    manager_enc = tf.train.CheckpointManager(ckpt_enc, 'checkpoints/' + config['name'] + '_enc', max_to_keep=20)
+    manager_enc = tf.train.CheckpointManager(ckpt_enc, 'checkpoints/post/' + config['name'] + '_enc', max_to_keep=20)
     ckpt_dec = tf.train.Checkpoint(model=model.decoder)
-    manager_dec = tf.train.CheckpointManager(ckpt_dec, 'checkpoints/' + config['name'] + '_dec', max_to_keep=20)
+    manager_dec = tf.train.CheckpointManager(ckpt_dec, 'checkpoints/post/' + config['name'] + '_dec', max_to_keep=20)
 
     # test()
     train()

calculate_error.py

@@ -1,6 +1,6 @@
 import os
 import sys
-from configs.config_abl_abs_skirt import config
+from configs.config_abl_all_post import config
 
 os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
 # import tensorflow as tf
@@ -59,13 +59,13 @@ if __name__ == '__main__':
     folder = config['test_dir']
 
     checkpoint_name = config['name']
-    iteration = '70000'
+    iteration = ''
 
     F = np.load("custom_utils/smpl_faces" + skirt_modifier + ".npy").astype(np.int32)
     garment_type_table = np.load('custom_data/garment_type_table_test.npy', allow_pickle=True).item()
 
 
-    output_folder = '/home/workspace/data/results/' + checkpoint_name + '-' + iteration
+    output_folder = '/home/workspace/data/results/post/' + checkpoint_name + '-' + iteration
 
 
     faces_folder = '/home/workspace/data/results/support/faces' + ('_skirt' if config['is_skirt'] else '')
@@ -73,7 +73,7 @@ if __name__ == '__main__':
     ids = [id.split('.')[0] for id in ids]
     faces = {id : np.load(faces_folder + '/' + id + '.npy') for id in ids }
 
-    masks_folder = '/home/workspace/data/dataset/test' + ('_skirt' if config['is_skirt'] else '') + '_ready/mesh_mask'
+    masks_folder = '/home/workspace/data/dataset/test' + ('_skirt' if config['is_skirt'] else '') + '_ready' + ('_2' if config['is_skirt'] else '') + '/mesh_mask'
     ids = os.listdir(masks_folder)
     ids = [id.split('.')[0] for id in ids]
     masks = {id : np.load(masks_folder + '/' + id + '.npy') for id in ids }
@@ -107,7 +107,7 @@ if __name__ == '__main__':
     for key in per_garment_errors.keys():
         print(key + ': ' + str(np.mean(per_garment_errors[key])))
 
-    with open('/home/workspace/data/results/errors/' + checkpoint_name + '-' + iteration + '.csv', mode='w') as f:
+    with open('/home/workspace/data/results/errors/post/' + checkpoint_name + '-' + iteration + '.csv', mode='w') as f:
         file_writer = csv.writer(f, delimiter=',', quotechar='"', quoting=csv.QUOTE_MINIMAL)
 
         for key in sorted(errors.keys()):

cape_convert.py

@@ -659,7 +659,7 @@ if __name__ == '__main__':
         # create_posed_body()
         # create_static()
         # create_unposed()
-        # create_posed()
+        create_posed()
         # verify()
         pass
 

cluster_poses.py

@@ -50,13 +50,13 @@ def sample_from_cluster(file, n_sample, output_file):
 
 
 if __name__ == '__main__':
-    folder = '/home/workspace/data/dataset/test_ready'
+    folder = '/home/workspace/data/dataset/test_skirt_ready_2'
 
-    is_skirt = False
+    is_skirt = True
     F = np.load("custom_utils/smpl_faces" + ("_skirt" if is_skirt == 1 else "") + ".npy")
 
-    # generate_clusters(folder, 20)
+    generate_clusters(folder, 20)
-    sample_from_cluster('custom_data/test_ready_pose_cluster_20.csv', 10, '/home/workspace/data/dataset/test_ready/frames_pose_cluster_20_sample_10.txt')
+    # sample_from_cluster('custom_data/test_skirt_ready_pose_cluster_20.csv', 10, '/home/workspace/data/dataset/test_skirt_ready/frames_pose_cluster_20_sample_10.txt')
 
 
     # samples = df.groupby('label', group_keys=False).apply(lambda x: x.sample())

compare_cape.py

@@ -0,0 +1,293 @@
+import os
+import sys
+
+os.environ["CUDA_VISIBLE_DEVICES"] = '6'
+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 tqdm
+import shutil
+import random
+import numpy as np
+import trimesh as tm
+import plotly.express as ex
+import plotly.graph_objects as go
+from custom_utils.custom_utils import _maskFaces
+from custom_utils.smpl import SMPL
+from plotly.subplots import make_subplots
+from scipy.spatial.transform import Rotation as R
+
+
+def plot_fixed(mesh, faces, colors):
+        fig = go.Figure()
+        
+        fig.add_trace(go.Mesh3d(
+                x=mesh[:,0],
+                y=mesh[:,2],
+                z=mesh[:,1],
+                i = faces[:,0],
+                j = faces[:,1],
+                k = faces[:,2],
+                text="",
+                vertexcolor=colors,
+                lighting={ 'ambient':0.8, 'diffuse':0.5, 'roughness':0.5 },
+                
+            ),)
+
+
+
+        fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1))
+        helper = dict(xaxis = dict(range=[-1,1], showbackground=False, showticklabels=False, showaxeslabels=False, showspikes=False, visible=False),
+                    yaxis = dict(range=[-1,1], showbackground=False, showticklabels=False, showaxeslabels=False, showspikes=False, visible=False),
+                    zaxis = dict(range=[-1,1], showbackground=False, showticklabels=False, showaxeslabels=False, showspikes=False, visible=False),)
+        fig.update_layout(scene = helper)
+        
+        camera = dict(
+            up=dict(x=0, y=0, z=1),
+            center=dict(x=0, y=0, z=-0.05),
+            eye=dict(x=0.3, y=1.0, z=0.05)
+        )
+
+        fig.update_layout(scene_camera=camera, width=800, height=1100, autosize=False,
+                          margin=dict(l=0, r=0, b=0, t=0, pad=0)
+                          )
+
+        return fig
+
+
+def plot_fixed_all(gt, cape, ours, error_cape, error_ours, faces, colors):
+        fig = make_subplots(rows=1, cols=3, start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(3)] ], horizontal_spacing=0)
+        
+        fig.add_trace(go.Mesh3d(
+                x=gt[:,0],
+                y=gt[:,2],
+                z=gt[:,1],
+                i = faces[:,0],
+                j = faces[:,1],
+                k = faces[:,2],
+                text='GT',
+                vertexcolor=colors,
+                lighting={ 'ambient':0.8, 'diffuse':0.5, 'roughness':0.5 },
+            ),
+            row=1, col=1)
+
+        fig.add_trace(go.Mesh3d(
+                x=cape[:,0],
+                y=cape[:,2],
+                z=cape[:,1],
+                i = faces[:,0],
+                j = faces[:,1],
+                k = faces[:,2],
+                text=str(error_cape),
+                vertexcolor=colors,
+                lighting={ 'ambient':0.8, 'diffuse':0.5, 'roughness':0.5 },
+            ),
+            row=1, col=2)
+
+
+        fig.add_trace(go.Mesh3d(
+                x=ours[:,0],
+                y=ours[:,2],
+                z=ours[:,1],
+                i = faces[:,0],
+                j = faces[:,1],
+                k = faces[:,2],
+                text=str(error_ours),
+                vertexcolor=colors,
+                lighting={ 'ambient':0.8, 'diffuse':0.5, 'roughness':0.5 },
+            ),
+            row=1, col=3)
+
+
+        fig.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))
+        helper = dict(xaxis = dict(range=[-1,1], showbackground=False, showticklabels=False, showaxeslabels=False, showspikes=False, visible=False),
+                    yaxis = dict(range=[-1,1], showbackground=False, showticklabels=False, showaxeslabels=False, showspikes=False, visible=False),
+                    zaxis = dict(range=[-1,1], showbackground=False, showticklabels=False, showaxeslabels=False, showspikes=False, visible=False),)
+        fig.update_layout(scene = helper,
+                          scene2 = helper,
+                          scene3 = helper)
+        
+        camera = dict(
+            up=dict(x=0, y=0, z=1),
+            center=dict(x=0, y=-0.3, z=-0.05),
+            eye=dict(x=0.5, y=1.2, z=0.05)
+        )
+
+        fig.update_layout(scene_camera=camera, width=2000, height=1100, autosize=False, margin=dict(l=0, r=0, b=0, t=0, pad=0),
+                          scene2_camera=camera,
+                          scene3_camera=camera)
+
+        return fig
+
+
+
+def custom_imshow(img):
+    fig = ex.imshow(np.flip(img, axis=0))
+    fig.update_layout(margin=dict(l=0, r=0, b=0, t=0, pad=0), width=img.shape[1], height=img.shape[0])
+    fig.update_xaxes(visible=False)
+    fig.update_yaxes(visible=False)
+
+    return fig
+
+
+def read_obj(path):
+    with open(path, 'r') as f:
+        obj = tm.exchange.obj.load_obj(f, skip_materials=True, maintain_order=True)
+    verts = obj['vertices']
+    return verts
+
+
+def read_frames_file(file):
+    with open(file, 'r') as f:
+        lines = f.readlines()
+    frames = ['_'.join(line.strip().split('.')[:-1]) for line in lines]
+
+    # frames_lookup = dict()
+
+    return frames
+
+
+def get_cape_errors(file, frames):
+    out = dict()
+    errors = np.load(file).mean(-1)
+
+    for frame, error in zip(frames, errors):
+        out[frame] = error
+
+    return out
+
+
+def get_errors(file):
+    out = dict()
+    with open(file, 'r') as f:
+        for line in f:
+            frame, error = line.split(' ')
+            out[frame] = float(error)
+    return out
+
+
+def read_gt(folder, frame):
+    split = frame.split('_')
+    file = '_'.join(split[:-1]) + '.' + split[-1] + '.npz'
+    data = np.load(os.path.join(folder, file))['mesh_posed']
+    return data
+
+
+def add_global_orientation(mesh, frame, folder, inverse=False):
+    split = frame.split('_')
+    file = '_'.join(split[:-1]) + '.' + split[-1] + '.npy'
+    pose_shape = np.load(os.path.join(folder, file))
+
+    rodr = smpl.run_rodrigues(pose_shape[:72][np.newaxis,:].astype(np.float32)).numpy()
+
+    r = R.from_matrix(rodr[0, 0, :, :])
+    out = r.apply(mesh, inverse=inverse)
+
+    return out
+
+
+if __name__ == '__main__':
+    folder_f = 'CAPE/results/CAPE-affineconv_nz18_pose24_clotype8_female/test_reconstruction_objs_dataset_female_4clotypes'
+    folder_m = 'CAPE/results/CAPE-affineconv_nz18_pose24_clotype8_male/test_reconstruction_objs_dataset_male_4clotypes'
+
+    frames_file_f = '/home/workspace/data/dataset/cape_processed/datasets/dataset_female_4clotypes_frames.txt'
+    frames_file_m = '/home/workspace/data/dataset/cape_processed/datasets/dataset_male_4clotypes_frames.txt'
+    
+    frames_f = read_frames_file(frames_file_f)
+    frames_f_to_i = { frame: i for i, frame in enumerate(frames_f) } 
+    errors_cape_f = get_cape_errors('CAPE/results/CAPE-affineconv_nz18_pose24_clotype8_female/errors.npy', frames_f)
+    errors_cape_f = dict(reversed(sorted(errors_cape_f.items(), key=lambda item: item[1])))
+    frames_m = read_frames_file(frames_file_m)
+    frames_m_to_i = { frame: i for i, frame in enumerate(frames_m) } 
+    errors_cape_m = get_cape_errors('CAPE/results/CAPE-affineconv_nz18_pose24_clotype8_male/errors.npy', frames_m)
+    errors_cape_m = dict(reversed(sorted(errors_cape_m.items(), key=lambda item: item[1])))
+
+    errors_f = get_errors('temp/error_abl_2_pyr_64_128_256_512_cape_f.txt')
+    errors_m = get_errors('temp/error_abl_2_pyr_64_128_256_512_cape_m.txt')
+
+    folder_pyr_f = '/home/workspace/data/results/cape/abl_2_pyr_64_128_256_512_cape_f'
+    folder_pyr_m = '/home/workspace/data/results/cape/abl_2_pyr_64_128_256_512_cape_m'
+
+    gt_folder_f = '/home/workspace/local-data/cape_processed/test/female/mesh_posed'
+    gt_folder_m = '/home/workspace/local-data/cape_processed/test/male/mesh_posed'
+
+    pose_folder_m = '/home/workspace/local-data/cape_processed/test/male/pose_shape'
+    pose_folder_f = '/home/workspace/local-data/cape_processed/test/female/pose_shape'
+
+
+    smpl = SMPL(orig_smpl=True)
+    # verts_idx = smpl.to_keep
+    clothing_verts_idx = np.load(os.path.join('CAPE', 'data', 'clothing_verts_idx.npy'))
+    cape_mask = np.zeros((6890), bool)
+    cape_mask[clothing_verts_idx] = True
+    cape_mask = cape_mask[smpl.to_keep]
+    J0 = np.load('custom_data/smpl_root_joint.npy')
+
+
+    _, faces = _maskFaces(np.zeros((3638, 3)), np.load('custom_utils/smpl_faces_orig.npy'), cape_mask)
+    colors = np.ones((3627, 3), np.int8) * 230
+
+    gender = 'male'
+
+    if gender == 'male':
+        folder_cape = folder_m
+        errors_cape = errors_cape_m
+        folder = folder_pyr_m
+        errors = errors_m
+        frames_to_i = frames_m_to_i
+        gt_folder = gt_folder_m
+        pose_folder = pose_folder_m
+    else:
+        folder_cape = folder_f
+        errors_cape = errors_cape_f
+        folder = folder_pyr_f
+        errors = errors_f
+        frames_to_i = frames_f_to_i
+        gt_folder = gt_folder_f
+        pose_folder = pose_folder_f
+
+    all_errors_1 = []
+    all_errors_2 = []
+
+    # for key in tqdm.tqdm(errors_cape.keys()):
+    count = 0
+    while True:
+        key = random.choice(list(errors_cape.keys()))
+        cape = read_obj(os.path.join(folder_cape, '{:0>4d}.obj'.format(frames_to_i[key])))[clothing_verts_idx, :]  - J0[np.newaxis, :]
+        cape = add_global_orientation(cape, key, pose_folder, inverse=True)
+        if not os.path.exists(os.path.join(folder, key + '_posed.npy')):
+            continue
+        ours = np.load(os.path.join(folder, key + '_posed.npy'))[cape_mask, :]
+        # ours = add_global_orientation(ours, key, pose_folder)
+        gt = read_gt(gt_folder, key)[cape_mask, :]
+        gt = add_global_orientation(gt, key, pose_folder, inverse=True)
+
+        # error_1 = np.sqrt(((gt - cape) ** 2).sum(-1)).mean()
+        # error_2 = np.sqrt(((gt - ours) ** 2).sum(-1)).mean()
+
+        # all_errors_1.append(error_1)
+        # all_errors_2.append(error_2)
+
+    # print(np.mean(all_errors_1))
+    # print(np.mean(all_errors_2))
+
+        fig = plot_fixed_all(gt, cape, ours, errors_cape[key], errors[key], faces, colors)
+        fig.write_image('temp/cape/' + key + '.png')
+
+        print(count)
+        count += 1
+
+    print()

compare_models.py

@@ -0,0 +1,273 @@
+import os
+import sys
+from configs.config_abl_all_skirt import config as config_a
+from configs.config_abl_extra_loss_skirt import config as config_b
+from configs.config_abl_norm_skirt import config as config_c    
+
+os.environ["CUDA_VISIBLE_DEVICES"] = '6'
+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 glob
+import time
+import argparse
+import wandb
+import tqdm
+import random
+import numpy as np
+import plotly.graph_objects as go
+import plotly.express as ex
+from PIL import Image
+from postprocessing import fix_collision
+
+if config_a['is_skirt']:
+    from losses import edge_loss_skirt as edge_loss
+else:
+    from losses import edge_loss
+
+
+
+from losses import collision_ratio, kl_loss #, concept_loss
+# from losses import collision_penalty as collision_penalty
+from losses import collision_loss as collision_penalty
+from custom_utils.custom_utils import batch_to_dict, load_model, plot_3d, UV2Mesh, get_summary_html, maskFacesBatch, _maskFaces, plot_3d_collision
+from dataset import DatasetFactory, StaticDataset
+from custom_utils.smpl import SMPL
+from models.cvae_model import CVAE
+from custom_utils.supersmpl_io import writeOBJ
+from plotly.subplots import make_subplots
+
+
+uv2mesh = UV2Mesh(res=256, is_skirt=config_a['is_skirt'], is_adaptive=False)
+
+
+
+def plot(mesh_gt, pred_a, pred_b, pred_c, body, mask):
+    fig = make_subplots(rows=1, cols=4, start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(4)] ])
+    
+    _, faces_garment = _maskFaces(np.zeros((14475, 3), np.float32), F, mask)
+    faces = np.concatenate((F, faces_garment + 14475), axis=0)
+    mesh_gt = np.concatenate((body, mesh_gt[mask, :]), axis=0)
+    pred_a = np.concatenate((body, pred_a[mask, :]), axis=0)
+    pred_b = np.concatenate((body, pred_b[mask, :]), axis=0)
+    pred_c = np.concatenate((body, pred_c[mask, :]), axis=0)
+
+    colors = np.concatenate((np.ones_like(body, dtype=np.int32)*100, np.ones_like(mesh_gt, dtype=np.int32)*230), axis=0)
+
+    fig.add_trace(go.Mesh3d(
+            x=mesh_gt[:,0],
+            y=mesh_gt[:,2],
+            z=mesh_gt[:,1],
+            i = faces[:,0],
+            j = faces[:,1],
+            k = faces[:,2],
+            text="Ground Truth",
+            vertexcolor=colors,
+        ),
+        row=1, col=1)
+
+    fig.add_trace(go.Mesh3d(
+            x=pred_a[:,0],
+            y=pred_a[:,2],
+            z=pred_a[:,1],
+            i = faces[:,0],
+            j = faces[:,1],
+            k = faces[:,2],
+            text="Pred A",
+            vertexcolor=colors,
+        ),
+        row=1, col=2)
+
+    fig.add_trace(go.Mesh3d(
+            x=pred_b[:,0],
+            y=pred_b[:,2],
+            z=pred_b[:,1],
+            i = faces[:,0],
+            j = faces[:,1],
+            k = faces[:,2],
+            text="Pred B",
+            vertexcolor=colors,
+        ),
+        row=1, col=3)
+
+    fig.add_trace(go.Mesh3d(
+            x=pred_c[:,0],
+            y=pred_c[:,2],
+            z=pred_c[:,1],
+            i = faces[:,0],
+            j = faces[:,1],
+            k = faces[:,2],
+            text="Pred C",
+            vertexcolor=colors,
+        ),
+        row=1, col=4)
+
+
+    fig.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],),)
+    fig.update_layout(scene = helper,
+                        scene2 = helper,
+                        scene3 = helper,
+                        scene4 = helper,
+                        )
+
+    return fig
+
+def step(model, config, batch):
+    if config['use_skinning']:
+        uv_orig = batch['uv_unposed']
+    else:
+        uv_orig =  batch['uv_posed']
+    
+    pred, _, _ = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False)
+
+    mesh_pred = uv2mesh(pred)
+
+    if config['use_skinning']:
+        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_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)
+    else:
+        mesh_pred_posed = mesh_pred
+        
+    # mesh_pred_posed_fixed = fix_collision(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
+
+    return mesh_pred_posed, None #mesh_pred_posed_fixed
+
+
+if __name__ == '__main__':
+    # init_wandb()
+    parser = argparse.ArgumentParser(description='')
+    parser.add_argument('--name', type=str)
+    args = vars(parser.parse_args())
+
+    skirt_modifier = '_skirt' if config_a['is_skirt'] else ''
+
+    folder = config_a['test_dir']
+    smpl = SMPL(use_one_hot=False)
+
+    model_a = CVAE(config_a)
+    model_b = CVAE(config_b)
+    model_c = CVAE(config_c)
+
+
+    F = np.load("custom_utils/smpl_faces" + skirt_modifier + ".npy").astype(np.int32)
+
+    iteration = '70000'
+
+    load_model(model_a, 'checkpoints/' + config_a['name'], iteration=iteration)
+    load_model(model_b, 'checkpoints/' + config_b['name'], iteration=iteration)
+    load_model(model_c, 'checkpoints/' + config_c['name'], iteration=iteration)
+    # load_model(model, 'checkpoints/' + 'abl_all_120k')
+
+    assert config_a['remove_global_orientation'] == config_b['remove_global_orientation'] == config_c['remove_global_orientation']
+
+    factory = DatasetFactory()
+    dataset_a = factory.get_dataset(folder=config_a['test_dir'], 
+                                       batch_size=config_a['batch_size'], 
+                                       signature=config_a['signature_dynamic'], 
+                                       resolutions=[config_a['resolution']], 
+                                       in_memory=False, 
+                                       do_repeat=False,
+                                       use_one_hot=config_a['use_one_hot'],
+                                       use_no_orient=config_a['remove_global_orientation'],
+                                       shuffle=False,
+                                       testing=True,
+                                    #    frames_file='frames_abl.txt',
+                                       )
+    static_dataset_a = StaticDataset(folder=config_a['test_dir'], 
+                                        signature=config_a['signature_static'], 
+                                        resolutions=[config_a['resolution']], 
+                                        in_memory=True)
+
+    dataset_b = factory.get_dataset(folder=config_b['test_dir'], 
+                                       batch_size=config_b['batch_size'], 
+                                       signature=config_b['signature_dynamic'], 
+                                       resolutions=[config_b['resolution']], 
+                                       in_memory=False, 
+                                       do_repeat=False,
+                                       use_one_hot=config_b['use_one_hot'],
+                                       use_no_orient=config_b['remove_global_orientation'],
+                                       shuffle=False,
+                                       testing=True,
+                                    #    frames_file='frames_abl.txt',
+                                       )
+    static_dataset_b = StaticDataset(folder=config_b['test_dir'], 
+                                        signature=config_b['signature_static'], 
+                                        resolutions=[config_b['resolution']], 
+                                        in_memory=True)
+
+    dataset_c = factory.get_dataset(folder=config_c['test_dir'], 
+                                       batch_size=config_c['batch_size'], 
+                                       signature=config_c['signature_dynamic'], 
+                                       resolutions=[config_c['resolution']], 
+                                       in_memory=False, 
+                                       do_repeat=False,
+                                       use_one_hot=config_c['use_one_hot'],
+                                       use_no_orient=config_c['remove_global_orientation'],
+                                       shuffle=False,
+                                       testing=True,
+                                    #    frames_file='frames_abl.txt',
+                                       )
+    static_dataset_c = StaticDataset(folder=config_c['test_dir'], 
+                                        signature=config_c['signature_static'], 
+                                        resolutions=[config_c['resolution']], 
+                                        in_memory=True)
+
+    iterator_a = iter(dataset_a)
+    iterator_b = iter(dataset_b)
+    iterator_c = iter(dataset_c)
+    
+
+    faces_folder = '/home/workspace/data/results/support/faces' + ('_skirt' if config_a['is_skirt'] else '')
+    ids = os.listdir(faces_folder)
+    ids = [id.split('.')[0] for id in ids]
+    faces = {id : np.load(faces_folder + '/' + id + '.npy') for id in ids }
+
+
+    while True:
+        batch_a = iterator_a.get_next()
+        batch_b = iterator_b.get_next()
+        batch_c = iterator_c.get_next()
+
+        batch_a = static_dataset_a.expand_batch(batch_to_dict(batch_a, config_a['signature_dynamic']))
+        batch_b = static_dataset_b.expand_batch(batch_to_dict(batch_b, config_b['signature_dynamic']))
+        batch_c = static_dataset_c.expand_batch(batch_to_dict(batch_c, config_c['signature_dynamic']))
+
+        mesh_pred_posed_a, mesh_pred_posed_fixed_a = step(model_a, config_a, batch_a)
+        mesh_pred_posed_b, mesh_pred_posed_fixed_b = step(model_b, config_b, batch_b)
+        mesh_pred_posed_c, mesh_pred_posed_fixed_c = step(model_c, config_c, batch_c)
+
+        mesh_pred_posed_a = mesh_pred_posed_a.numpy()
+        mesh_pred_posed_b = mesh_pred_posed_b.numpy()
+        mesh_pred_posed_c = mesh_pred_posed_c.numpy()
+        mask = batch_a['mesh_mask'].numpy().astype(bool)
+        gt = batch_a['mesh_posed'].numpy()
+        body = batch_a['mesh_body_posed'].numpy()
+
+        for i in range(mask.shape[0]):
+            id, frame_nr = batch_a['id'][i].numpy().decode("utf-8"), batch_a['frame_nr'][i].numpy().decode("utf-8")
+            plot(gt[i], mesh_pred_posed_a[i], mesh_pred_posed_b[i], mesh_pred_posed_c[i], body[i], mask[i]).show()
+
+            print()
+

configs/config_abl_2_pyr_32.py

@@ -1,6 +1,6 @@
 config = dict()
 
-config['gpu'] = '1'
+config['gpu'] = '6'
 config['name'] = 'abl_2_pyr_32'
 config['notes'] = 'fixed noise handling'
 config['tags'] = ['pyr_non_skirt_2']

configs/config_abl_2_pyr_32_64.py

@@ -1,6 +1,6 @@
 config = dict()
 
-config['gpu'] = '4'
+config['gpu'] = '6'
 config['name'] = 'abl_2_pyr_32_64'
 config['notes'] = 'fixed noise handling'
 config['tags'] = ['pyr_non_skirt_2']

configs/config_abl_2_pyr_32_64_128.py

@@ -0,0 +1,186 @@
+config = dict()
+
+config['gpu'] = '6'
+config['name'] = 'abl_2_pyr_32_64_128'
+config['notes'] = 'fixed noise handling'
+config['tags'] = ['pyr_non_skirt_2']
+
+# DON'T FORGET TO RESET SSIM LOSS USAGE
+# DON'T FORGET TO CHECK MODEL LOADING
+# DON'T FORGET TO CHECK IF STORE IN MEMORY
+# DON'T FORGET TO CHECK REG LOSS COEFF
+# DON'T FORGET TO LOAD FINAL UPSCALER UNLESS HAVE LEVEL 512
+# DON'T FORGET TO CHECK IF SKIRT
+# DON'T FORGET TO CHECK IF CAPE
+# DON'T FORGET TO CHECK IF USING DOWNSCALE NONE OR _2
+# DON'T FORGET TO CHECK UPSCALING and ENCODER CHECKPOINTS
+
+config['use_skinning'] =              True # and deshaping
+config['remove_global_orientation'] = True
+config['use_extra_losses'] =          True # edge, collision, 3d mesh
+config['use_offsets'] =               False
+config['use_norm_loss'] =             False
+
+config['pretrained_model'] = ''
+config['pretrained_model_iteration'] = ''
+
+config['pyramid'] = dict()
+config['pyramid']['resolutions'] = [32, 64, 128]
+config['pyramid']['pretrained_pyramid'] = 'checkpoints/abl_2_pyr_32_64'
+config['pyramid']['pretrained_pyramid_iteration'] = '100000'
+config['pyramid']['final_upscale_enhancer'] = 'checkpoints/128_512_3'
+
+config['train_dir'] = '/home/workspace/local-data/train'
+config['valid_dir'] = '/home/workspace/local-data/valid'
+config['test_dir'] = '/home/workspace/data/dataset/test_ready'
+config['batch_size'] = 4
+config['is_skirt'] = False
+config['is_cape'] = False
+
+config['dataset_in_memory'] = False
+config['resolution'] = None
+config['latent_dim'] = 128
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['lr'] = 0.0001
+config['starting_balancing_coeff'] = 0.0001
+config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
+config['avg_garment_path'] = ''
+
+config['signature_dynamic'] = { 'id': (), 'frame_nr': (),
+                          'pose_shape': (72 + 11,),
+                          'uv_body_posed': (512, 256, 3),
+                          'mesh_body_posed': (config['smpl_res'], 3,),
+                          'mesh_posed': (config['smpl_res'], 3,),
+                          'mesh_unposed': (config['smpl_res'], 3,),
+                          'uv_unposed': (512, 256, 3),
+                          'uv_curvature': (512, 256, 7),
+                        }
+config['signature_static'] = { 'uv_static': (512, 256, 3),
+                              'mesh_mask': (config['smpl_res'],),
+                              'uv_mask': (512, 256,),
+                              }
+
+config['pyramid']['is_label'] = False
+config['pyramid']['latent_dim'] = config['latent_dim']
+config['pyramid']['add_noise'] = True
+config['pyramid']['noise_amount'] = 0.001
+
+config['pyramid']['loss_weights'] = [1,1,1,1,1]
+
+for res in config['pyramid']['resolutions']:
+  if res != 512:
+      res_modifier = '_' + str(res)
+      config['signature_dynamic']['uv_body_posed' + res_modifier] = (res, res//2, 3)
+      config['signature_dynamic']['uv_unposed' + res_modifier] = (res, res//2, 3)
+      config['signature_dynamic']['uv_curvature' + res_modifier] = (res, res//2, 7)
+      config['signature_static']['uv_static' + res_modifier] = (res, res//2, 3)
+      config['signature_static']['uv_mask' + res_modifier] = (res, res//2)
+
+
+config['pyramid']['static_model'] = 'checkpoints/static_1'
+config['pyramid']['body_model'] = 'checkpoints/body_1'
+config['pyramid']['curvature_model'] = 'checkpoints/curv_1'
+config['pyramid']['freeze_encoders'] = True
+
+config['pyramid']['levels'] = []
+for res in config['pyramid']['resolutions']:
+  if res == 512:
+    level = dict()
+    level['resolution'] = 512
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 7
+    level['encoder_depths'] = [3, 3, 9, 3]
+    level['encoder_dims'] = [32, 64, 128, 256]
+    level['decoder_input'] = (16, 8, 80)
+    level['decoder_depths'] = [3, 3, 9, 3]
+    level['decoder_dims'] = [256, 128, 64, 16]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = ''
+    config['pyramid']['levels'].append(level)
+  elif res == 256:
+    level = dict()
+    level['resolution'] = 256
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 7
+    level['encoder_depths'] = [9, 9, 9, 9]
+    level['encoder_dims'] = [32, 64, 128, 256]
+    level['decoder_input'] = (8, 4, 128)
+    level['decoder_depths'] = [9, 9, 9, 9]
+    level['decoder_dims'] = [256, 128, 64, 16]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = ''#'checkpoints/up_3_d256_512_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  elif res == 128:
+    level = dict()
+    level['resolution'] = 128
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 5
+    level['encoder_depths'] = [16, 9, 9, 9]
+    level['encoder_dims'] = [64, 128, 128, 256]
+    level['decoder_input'] = (4, 2, 128)
+    level['decoder_depths'] = [16, 9, 9, 9]
+    level['decoder_dims'] = [256, 128, 128, 64]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = ''#'checkpoints/up_3_d128_256_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  elif res == 64:
+    level = dict()
+    level['resolution'] = 64
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 5
+    level['encoder_depths'] = [9, 9, 9]
+    level['encoder_dims'] = [64, 128, 256]
+    level['decoder_input'] = (4, 2, 256)
+    level['decoder_depths'] = [9, 9, 9]
+    level['decoder_dims'] = [256, 128, 64]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = 'checkpoints/64_512_2'#checkpoints/up_3_d64_128_2'
+    level['upscale_enhancer_freeze'] = False
+    config['pyramid']['levels'].append(level)
+  elif res == 32:
+    level = dict()
+    level['resolution'] = 32
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 3
+    level['encoder_depths'] = [9, 9]
+    level['encoder_dims'] = [256, 256]
+    level['decoder_input'] = (4, 2, 256)
+    level['decoder_depths'] = [9, 9]
+    level['decoder_dims'] = [256, 256]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = 'checkpoints/32_512_4'#'checkpoints/up_3_d32_64_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  else:
+    raise NotImplementedError
+
+
+
+config['cond_ae'] = dict()
+config['cond_ae']['is_label'] = False
+config['cond_ae']['upscale_output'] = False
+config['cond_ae']['latent_dim'] = config['latent_dim']
+config['cond_ae']['kernel_size'] = 7
+config['cond_ae']['bottleneck_size'] = 64
+config['cond_ae']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae']['encoder_dims'] = [32, 64, 128, 256, config['cond_ae']['bottleneck_size']]
+config['cond_ae']['decoder_input'] = (8, 4, config['cond_ae']['bottleneck_size'])
+config['cond_ae']['decoder_depths'] = [3, 3, 9, 3, 3]
+config['cond_ae']['decoder_dims'] = [config['cond_ae']['bottleneck_size'], 256, 128, 64, 16]
+
+config['cond_ae_body'] = dict()
+config['cond_ae_body']['latent_dim'] = config['latent_dim']
+config['cond_ae_body']['kernel_size'] = 7
+config['cond_ae_body']['upscale_output'] = True
+config['cond_ae_body']['bottleneck_size'] = 512
+config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3, 3]
+config['cond_ae_body']['encoder_dims'] = [32, 64, 128, 256, 256, config['cond_ae_body']['bottleneck_size']]
+config['cond_ae_body']['decoder_input'] = (4, 2, config['cond_ae_body']['bottleneck_size'])
+config['cond_ae_body']['decoder_depths'] = [3, 3, 9, 3]
+config['cond_ae_body']['decoder_dims'] = [config['cond_ae_body']['bottleneck_size'], 128, 64, 16]

configs/config_abl_2_pyr_32_64_128_256.py

@@ -0,0 +1,186 @@
+config = dict()
+
+config['gpu'] = '6'
+config['name'] = 'abl_2_pyr_32_64_128_256'
+config['notes'] = 'fixed noise handling'
+config['tags'] = ['pyr_non_skirt_2']
+
+# DON'T FORGET TO RESET SSIM LOSS USAGE
+# DON'T FORGET TO CHECK MODEL LOADING
+# DON'T FORGET TO CHECK IF STORE IN MEMORY
+# DON'T FORGET TO CHECK REG LOSS COEFF
+# DON'T FORGET TO LOAD FINAL UPSCALER UNLESS HAVE LEVEL 512
+# DON'T FORGET TO CHECK IF SKIRT
+# DON'T FORGET TO CHECK IF CAPE
+# DON'T FORGET TO CHECK IF USING DOWNSCALE NONE OR _2
+# DON'T FORGET TO CHECK UPSCALING and ENCODER CHECKPOINTS
+
+config['use_skinning'] =              True # and deshaping
+config['remove_global_orientation'] = True
+config['use_extra_losses'] =          True # edge, collision, 3d mesh
+config['use_offsets'] =               False
+config['use_norm_loss'] =             False
+
+config['pretrained_model'] = ''
+config['pretrained_model_iteration'] = ''
+
+config['pyramid'] = dict()
+config['pyramid']['resolutions'] = [32, 64, 128, 256]
+config['pyramid']['pretrained_pyramid'] = 'checkpoints/abl_2_pyr_32_64_128'
+config['pyramid']['pretrained_pyramid_iteration'] = '80000'
+config['pyramid']['final_upscale_enhancer'] = 'checkpoints/256_512'
+
+config['train_dir'] = '/home/workspace/local-data/train'
+config['valid_dir'] = '/home/workspace/local-data/valid'
+config['test_dir'] = '/home/workspace/data/dataset/test_ready'
+config['batch_size'] = 4
+config['is_skirt'] = False
+config['is_cape'] = False
+
+config['dataset_in_memory'] = False
+config['resolution'] = None
+config['latent_dim'] = 128
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['lr'] = 0.0001
+config['starting_balancing_coeff'] = 0.0001
+config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
+config['avg_garment_path'] = ''
+
+config['signature_dynamic'] = { 'id': (), 'frame_nr': (),
+                          'pose_shape': (72 + 11,),
+                          'uv_body_posed': (512, 256, 3),
+                          'mesh_body_posed': (config['smpl_res'], 3,),
+                          'mesh_posed': (config['smpl_res'], 3,),
+                          'mesh_unposed': (config['smpl_res'], 3,),
+                          'uv_unposed': (512, 256, 3),
+                          'uv_curvature': (512, 256, 7),
+                        }
+config['signature_static'] = { 'uv_static': (512, 256, 3),
+                              'mesh_mask': (config['smpl_res'],),
+                              'uv_mask': (512, 256,),
+                              }
+
+config['pyramid']['is_label'] = False
+config['pyramid']['latent_dim'] = config['latent_dim']
+config['pyramid']['add_noise'] = True
+config['pyramid']['noise_amount'] = 0.001
+
+config['pyramid']['loss_weights'] = [1,1,1,1,1]
+
+for res in config['pyramid']['resolutions']:
+  if res != 512:
+      res_modifier = '_' + str(res)
+      config['signature_dynamic']['uv_body_posed' + res_modifier] = (res, res//2, 3)
+      config['signature_dynamic']['uv_unposed' + res_modifier] = (res, res//2, 3)
+      config['signature_dynamic']['uv_curvature' + res_modifier] = (res, res//2, 7)
+      config['signature_static']['uv_static' + res_modifier] = (res, res//2, 3)
+      config['signature_static']['uv_mask' + res_modifier] = (res, res//2)
+
+
+config['pyramid']['static_model'] = 'checkpoints/static_1'
+config['pyramid']['body_model'] = 'checkpoints/body_1'
+config['pyramid']['curvature_model'] = 'checkpoints/curv_1'
+config['pyramid']['freeze_encoders'] = True
+
+config['pyramid']['levels'] = []
+for res in config['pyramid']['resolutions']:
+  if res == 512:
+    level = dict()
+    level['resolution'] = 512
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 7
+    level['encoder_depths'] = [3, 3, 9, 3]
+    level['encoder_dims'] = [32, 64, 128, 256]
+    level['decoder_input'] = (16, 8, 80)
+    level['decoder_depths'] = [3, 3, 9, 3]
+    level['decoder_dims'] = [256, 128, 64, 16]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = ''
+    config['pyramid']['levels'].append(level)
+  elif res == 256:
+    level = dict()
+    level['resolution'] = 256
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 7
+    level['encoder_depths'] = [9, 9, 9, 9]
+    level['encoder_dims'] = [32, 64, 128, 256]
+    level['decoder_input'] = (8, 4, 128)
+    level['decoder_depths'] = [9, 9, 9, 9]
+    level['decoder_dims'] = [256, 128, 64, 16]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = ''#'checkpoints/up_3_d256_512_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  elif res == 128:
+    level = dict()
+    level['resolution'] = 128
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 5
+    level['encoder_depths'] = [16, 9, 9, 9]
+    level['encoder_dims'] = [64, 128, 128, 256]
+    level['decoder_input'] = (4, 2, 128)
+    level['decoder_depths'] = [16, 9, 9, 9]
+    level['decoder_dims'] = [256, 128, 128, 64]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = 'checkpoints/128_512_3'#'checkpoints/up_3_d128_256_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  elif res == 64:
+    level = dict()
+    level['resolution'] = 64
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 5
+    level['encoder_depths'] = [9, 9, 9]
+    level['encoder_dims'] = [64, 128, 256]
+    level['decoder_input'] = (4, 2, 256)
+    level['decoder_depths'] = [9, 9, 9]
+    level['decoder_dims'] = [256, 128, 64]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = 'checkpoints/64_512_2'#checkpoints/up_3_d64_128_2'
+    level['upscale_enhancer_freeze'] = False
+    config['pyramid']['levels'].append(level)
+  elif res == 32:
+    level = dict()
+    level['resolution'] = 32
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 3
+    level['encoder_depths'] = [9, 9]
+    level['encoder_dims'] = [256, 256]
+    level['decoder_input'] = (4, 2, 256)
+    level['decoder_depths'] = [9, 9]
+    level['decoder_dims'] = [256, 256]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = 'checkpoints/32_512_4'#'checkpoints/up_3_d32_64_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  else:
+    raise NotImplementedError
+
+
+
+config['cond_ae'] = dict()
+config['cond_ae']['is_label'] = False
+config['cond_ae']['upscale_output'] = False
+config['cond_ae']['latent_dim'] = config['latent_dim']
+config['cond_ae']['kernel_size'] = 7
+config['cond_ae']['bottleneck_size'] = 64
+config['cond_ae']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae']['encoder_dims'] = [32, 64, 128, 256, config['cond_ae']['bottleneck_size']]
+config['cond_ae']['decoder_input'] = (8, 4, config['cond_ae']['bottleneck_size'])
+config['cond_ae']['decoder_depths'] = [3, 3, 9, 3, 3]
+config['cond_ae']['decoder_dims'] = [config['cond_ae']['bottleneck_size'], 256, 128, 64, 16]
+
+config['cond_ae_body'] = dict()
+config['cond_ae_body']['latent_dim'] = config['latent_dim']
+config['cond_ae_body']['kernel_size'] = 7
+config['cond_ae_body']['upscale_output'] = True
+config['cond_ae_body']['bottleneck_size'] = 512
+config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3, 3]
+config['cond_ae_body']['encoder_dims'] = [32, 64, 128, 256, 256, config['cond_ae_body']['bottleneck_size']]
+config['cond_ae_body']['decoder_input'] = (4, 2, config['cond_ae_body']['bottleneck_size'])
+config['cond_ae_body']['decoder_depths'] = [3, 3, 9, 3]
+config['cond_ae_body']['decoder_dims'] = [config['cond_ae_body']['bottleneck_size'], 128, 64, 16]

configs/config_abl_2_pyr_32_64_128_256_512.py

@@ -0,0 +1,191 @@
+config = dict()
+
+config['gpu'] = '6'
+config['name'] = 'abl_2_pyr_32_64_128_256_512'
+config['notes'] = 'fixed noise handling'
+config['tags'] = ['pyr_non_skirt_2']
+
+# DON'T FORGET TO RESET SSIM LOSS USAGE
+# DON'T FORGET TO CHECK MODEL LOADING
+# DON'T FORGET TO CHECK IF STORE IN MEMORY
+# DON'T FORGET TO CHECK REG LOSS COEFF
+# DON'T FORGET TO LOAD FINAL UPSCALER UNLESS HAVE LEVEL 512
+# DON'T FORGET TO CHECK IF SKIRT
+# DON'T FORGET TO CHECK IF CAPE
+# DON'T FORGET TO CHECK IF USING DOWNSCALE NONE OR _2
+# DON'T FORGET TO CHECK UPSCALING and ENCODER CHECKPOINTS
+
+config['use_skinning'] =              True # and deshaping
+config['remove_global_orientation'] = True
+config['use_extra_losses'] =          True # edge, collision, 3d mesh
+config['use_offsets'] =               False
+config['use_norm_loss'] =             False
+
+config['pretrained_model'] = 'checkpoints/abl_2_pyr_32_64_128_256_512'
+config['pretrained_model_iteration'] = '50000'
+
+config['pyramid'] = dict()
+config['pyramid']['resolutions'] = [32, 64, 128, 256, 512]
+config['pyramid']['pretrained_pyramid'] = 'checkpoints/abl_2_pyr_32_64_128_256'
+config['pyramid']['pretrained_pyramid_iteration'] = '70000'
+config['pyramid']['final_upscale_enhancer'] = ''
+config['pyramid']['upscale_enhancer_kernel_size'] = 3
+
+config['train_dir'] = '/home/workspace/local-data/train'
+config['valid_dir'] = '/home/workspace/local-data/valid'
+config['test_dir'] = '/home/workspace/data/dataset/test_ready'
+config['batch_size'] = 4
+config['is_skirt'] = False
+config['is_cape'] = False
+
+config['dataset_in_memory'] = False
+config['resolution'] = None
+config['latent_dim'] = 128
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['lr'] = 0.0001
+config['starting_balancing_coeff'] = 0.0001
+config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
+config['avg_garment_path'] = ''
+
+config['signature_dynamic'] = { 'id': (), 'frame_nr': (),
+                          'pose_shape': (72 + 11,),
+                          'uv_body_posed': (512, 256, 3),
+                          'mesh_body_posed': (config['smpl_res'], 3,),
+                          'mesh_posed': (config['smpl_res'], 3,),
+                          'mesh_unposed': (config['smpl_res'], 3,),
+                          'uv_unposed': (512, 256, 3),
+                          'uv_curvature': (512, 256, 7),
+                        }
+config['signature_static'] = { 'uv_static': (512, 256, 3),
+                              'mesh_mask': (config['smpl_res'],),
+                              'uv_mask': (512, 256,),
+                              }
+
+config['pyramid']['is_label'] = False
+config['pyramid']['latent_dim'] = config['latent_dim']
+config['pyramid']['add_noise'] = True
+config['pyramid']['noise_amount'] = 0.001
+
+config['pyramid']['loss_weights'] = [1,1,1,1,1]
+
+for res in config['pyramid']['resolutions']:
+  if res != 512:
+      res_modifier = '_' + str(res)
+      config['signature_dynamic']['uv_body_posed' + res_modifier] = (res, res//2, 3)
+      config['signature_dynamic']['uv_unposed' + res_modifier] = (res, res//2, 3)
+      config['signature_dynamic']['uv_curvature' + res_modifier] = (res, res//2, 7)
+      config['signature_static']['uv_static' + res_modifier] = (res, res//2, 3)
+      config['signature_static']['uv_mask' + res_modifier] = (res, res//2)
+
+
+config['pyramid']['static_model'] = 'checkpoints/static_1'
+config['pyramid']['body_model'] = 'checkpoints/body_1'
+config['pyramid']['curvature_model'] = 'checkpoints/curv_1'
+config['pyramid']['freeze_encoders'] = True
+
+config['pyramid']['levels'] = []
+for res in config['pyramid']['resolutions']:
+  if res == 512:
+    level = dict()
+    level['resolution'] = 512
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 7
+    level['encoder_depths'] = [3, 3, 9, 3]
+    level['encoder_dims'] = [32, 64, 128, 256]
+    level['decoder_input'] = (16, 8, 80)
+    level['decoder_depths'] = [3, 3, 9, 3]
+    level['decoder_dims'] = [256, 128, 64, 16]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = ''
+    config['pyramid']['levels'].append(level)
+  elif res == 256:
+    level = dict()
+    level['resolution'] = 256
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 7
+    level['encoder_depths'] = [9, 9, 9, 9]
+    level['encoder_dims'] = [32, 64, 128, 256]
+    level['decoder_input'] = (8, 4, 128)
+    level['decoder_depths'] = [9, 9, 9, 9]
+    level['decoder_dims'] = [256, 128, 64, 16]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = 'checkpoints/256_512'#'checkpoints/up_3_d256_512_2'
+    level['upscale_enhancer_freeze'] = True
+    level['upscale_enhancer_kernel_size'] = 3
+    config['pyramid']['levels'].append(level)
+  elif res == 128:
+    level = dict()
+    level['resolution'] = 128
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 5
+    level['encoder_depths'] = [16, 9, 9, 9]
+    level['encoder_dims'] = [64, 128, 128, 256]
+    level['decoder_input'] = (4, 2, 128)
+    level['decoder_depths'] = [16, 9, 9, 9]
+    level['decoder_dims'] = [256, 128, 128, 64]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = 'checkpoints/128_512_3'#'checkpoints/up_3_d128_256_2'
+    level['upscale_enhancer_freeze'] = True
+    level['upscale_enhancer_kernel_size'] = 3
+    config['pyramid']['levels'].append(level)
+  elif res == 64:
+    level = dict()
+    level['resolution'] = 64
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 5
+    level['encoder_depths'] = [9, 9, 9]
+    level['encoder_dims'] = [64, 128, 256]
+    level['decoder_input'] = (4, 2, 256)
+    level['decoder_depths'] = [9, 9, 9]
+    level['decoder_dims'] = [256, 128, 64]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = 'checkpoints/64_512_2'#checkpoints/up_3_d64_128_2'
+    level['upscale_enhancer_freeze'] = False
+    level['upscale_enhancer_kernel_size'] = 3
+    config['pyramid']['levels'].append(level)
+  elif res == 32:
+    level = dict()
+    level['resolution'] = 32
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 3
+    level['encoder_depths'] = [9, 9]
+    level['encoder_dims'] = [256, 256]
+    level['decoder_input'] = (4, 2, 256)
+    level['decoder_depths'] = [9, 9]
+    level['decoder_dims'] = [256, 256]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = 'checkpoints/32_512_4'#'checkpoints/up_3_d32_64_2'
+    level['upscale_enhancer_freeze'] = True
+    level['upscale_enhancer_kernel_size'] = 3
+    config['pyramid']['levels'].append(level)
+  else:
+    raise NotImplementedError
+
+
+
+config['cond_ae'] = dict()
+config['cond_ae']['is_label'] = False
+config['cond_ae']['upscale_output'] = False
+config['cond_ae']['latent_dim'] = config['latent_dim']
+config['cond_ae']['kernel_size'] = 7
+config['cond_ae']['bottleneck_size'] = 64
+config['cond_ae']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae']['encoder_dims'] = [32, 64, 128, 256, config['cond_ae']['bottleneck_size']]
+config['cond_ae']['decoder_input'] = (8, 4, config['cond_ae']['bottleneck_size'])
+config['cond_ae']['decoder_depths'] = [3, 3, 9, 3, 3]
+config['cond_ae']['decoder_dims'] = [config['cond_ae']['bottleneck_size'], 256, 128, 64, 16]
+
+config['cond_ae_body'] = dict()
+config['cond_ae_body']['latent_dim'] = config['latent_dim']
+config['cond_ae_body']['kernel_size'] = 7
+config['cond_ae_body']['upscale_output'] = True
+config['cond_ae_body']['bottleneck_size'] = 512
+config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3, 3]
+config['cond_ae_body']['encoder_dims'] = [32, 64, 128, 256, 256, config['cond_ae_body']['bottleneck_size']]
+config['cond_ae_body']['decoder_input'] = (4, 2, config['cond_ae_body']['bottleneck_size'])
+config['cond_ae_body']['decoder_depths'] = [3, 3, 9, 3]
+config['cond_ae_body']['decoder_dims'] = [config['cond_ae_body']['bottleneck_size'], 128, 64, 16]

configs/config_abl_2_pyr_32_64_128_256_512_norm.py

@@ -0,0 +1,186 @@
+config = dict()
+
+config['gpu'] = '6'
+config['name'] = 'abl_2_pyr_32_64_128_256_512_norm'
+config['notes'] = 'fixed noise handling'
+config['tags'] = ['pyr_non_skirt_2']
+
+# DON'T FORGET TO RESET SSIM LOSS USAGE
+# DON'T FORGET TO CHECK MODEL LOADING
+# DON'T FORGET TO CHECK IF STORE IN MEMORY
+# DON'T FORGET TO CHECK REG LOSS COEFF
+# DON'T FORGET TO LOAD FINAL UPSCALER UNLESS HAVE LEVEL 512
+# DON'T FORGET TO CHECK IF SKIRT
+# DON'T FORGET TO CHECK IF CAPE
+# DON'T FORGET TO CHECK IF USING DOWNSCALE NONE OR _2
+# DON'T FORGET TO CHECK UPSCALING and ENCODER CHECKPOINTS
+
+config['use_skinning'] =              True # and deshaping
+config['remove_global_orientation'] = True
+config['use_extra_losses'] =          True # edge, collision, 3d mesh
+config['use_offsets'] =               False
+config['use_norm_loss'] =             True
+
+config['pretrained_model'] = ''
+config['pretrained_model_iteration'] = ''
+
+config['pyramid'] = dict()
+config['pyramid']['resolutions'] = [32, 64, 128, 256, 512]
+config['pyramid']['pretrained_pyramid'] = 'checkpoints/abl_2_pyr_32_64_128_256'
+config['pyramid']['pretrained_pyramid_iteration'] = '70000'
+config['pyramid']['final_upscale_enhancer'] = ''
+
+config['train_dir'] = '/home/workspace/local-data/train'
+config['valid_dir'] = '/home/workspace/local-data/valid'
+config['test_dir'] = '/home/workspace/data/dataset/test_ready'
+config['batch_size'] = 4
+config['is_skirt'] = False
+config['is_cape'] = False
+
+config['dataset_in_memory'] = False
+config['resolution'] = None
+config['latent_dim'] = 128
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['lr'] = 0.0001
+config['starting_balancing_coeff'] = 0.0001
+config['reg_loss_coeff'] = 1.0 # set to 0 if don't want to use edge+collision
+config['avg_garment_path'] = ''
+
+config['signature_dynamic'] = { 'id': (), 'frame_nr': (),
+                          'pose_shape': (72 + 11,),
+                          'uv_body_posed': (512, 256, 3),
+                          'mesh_body_posed': (config['smpl_res'], 3,),
+                          'mesh_posed': (config['smpl_res'], 3,),
+                          'mesh_unposed': (config['smpl_res'], 3,),
+                          'uv_unposed': (512, 256, 3),
+                          'uv_curvature': (512, 256, 7),
+                        }
+config['signature_static'] = { 'uv_static': (512, 256, 3),
+                              'mesh_mask': (config['smpl_res'],),
+                              'uv_mask': (512, 256,),
+                              }
+
+config['pyramid']['is_label'] = False
+config['pyramid']['latent_dim'] = config['latent_dim']
+config['pyramid']['add_noise'] = True
+config['pyramid']['noise_amount'] = 0.001
+
+config['pyramid']['loss_weights'] = [1,1,1,1,1]
+
+for res in config['pyramid']['resolutions']:
+  if res != 512:
+      res_modifier = '_' + str(res)
+      config['signature_dynamic']['uv_body_posed' + res_modifier] = (res, res//2, 3)
+      config['signature_dynamic']['uv_unposed' + res_modifier] = (res, res//2, 3)
+      config['signature_dynamic']['uv_curvature' + res_modifier] = (res, res//2, 7)
+      config['signature_static']['uv_static' + res_modifier] = (res, res//2, 3)
+      config['signature_static']['uv_mask' + res_modifier] = (res, res//2)
+
+
+config['pyramid']['static_model'] = 'checkpoints/static_1'
+config['pyramid']['body_model'] = 'checkpoints/body_1'
+config['pyramid']['curvature_model'] = 'checkpoints/curv_1'
+config['pyramid']['freeze_encoders'] = True
+
+config['pyramid']['levels'] = []
+for res in config['pyramid']['resolutions']:
+  if res == 512:
+    level = dict()
+    level['resolution'] = 512
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 7
+    level['encoder_depths'] = [3, 3, 9, 3]
+    level['encoder_dims'] = [32, 64, 128, 256]
+    level['decoder_input'] = (16, 8, 80)
+    level['decoder_depths'] = [3, 3, 9, 3]
+    level['decoder_dims'] = [256, 128, 64, 16]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = ''
+    config['pyramid']['levels'].append(level)
+  elif res == 256:
+    level = dict()
+    level['resolution'] = 256
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 7
+    level['encoder_depths'] = [9, 9, 9, 9]
+    level['encoder_dims'] = [32, 64, 128, 256]
+    level['decoder_input'] = (8, 4, 128)
+    level['decoder_depths'] = [9, 9, 9, 9]
+    level['decoder_dims'] = [256, 128, 64, 16]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = 'checkpoints/256_512'#'checkpoints/up_3_d256_512_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  elif res == 128:
+    level = dict()
+    level['resolution'] = 128
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 5
+    level['encoder_depths'] = [16, 9, 9, 9]
+    level['encoder_dims'] = [64, 128, 128, 256]
+    level['decoder_input'] = (4, 2, 128)
+    level['decoder_depths'] = [16, 9, 9, 9]
+    level['decoder_dims'] = [256, 128, 128, 64]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = 'checkpoints/128_512_3'#'checkpoints/up_3_d128_256_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  elif res == 64:
+    level = dict()
+    level['resolution'] = 64
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 5
+    level['encoder_depths'] = [9, 9, 9]
+    level['encoder_dims'] = [64, 128, 256]
+    level['decoder_input'] = (4, 2, 256)
+    level['decoder_depths'] = [9, 9, 9]
+    level['decoder_dims'] = [256, 128, 64]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = 'checkpoints/64_512_2'#checkpoints/up_3_d64_128_2'
+    level['upscale_enhancer_freeze'] = False
+    config['pyramid']['levels'].append(level)
+  elif res == 32:
+    level = dict()
+    level['resolution'] = 32
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 3
+    level['encoder_depths'] = [9, 9]
+    level['encoder_dims'] = [256, 256]
+    level['decoder_input'] = (4, 2, 256)
+    level['decoder_depths'] = [9, 9]
+    level['decoder_dims'] = [256, 256]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = 'checkpoints/32_512_4'#'checkpoints/up_3_d32_64_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  else:
+    raise NotImplementedError
+
+
+
+config['cond_ae'] = dict()
+config['cond_ae']['is_label'] = False
+config['cond_ae']['upscale_output'] = False
+config['cond_ae']['latent_dim'] = config['latent_dim']
+config['cond_ae']['kernel_size'] = 7
+config['cond_ae']['bottleneck_size'] = 64
+config['cond_ae']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae']['encoder_dims'] = [32, 64, 128, 256, config['cond_ae']['bottleneck_size']]
+config['cond_ae']['decoder_input'] = (8, 4, config['cond_ae']['bottleneck_size'])
+config['cond_ae']['decoder_depths'] = [3, 3, 9, 3, 3]
+config['cond_ae']['decoder_dims'] = [config['cond_ae']['bottleneck_size'], 256, 128, 64, 16]
+
+config['cond_ae_body'] = dict()
+config['cond_ae_body']['latent_dim'] = config['latent_dim']
+config['cond_ae_body']['kernel_size'] = 7
+config['cond_ae_body']['upscale_output'] = True
+config['cond_ae_body']['bottleneck_size'] = 512
+config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3, 3]
+config['cond_ae_body']['encoder_dims'] = [32, 64, 128, 256, 256, config['cond_ae_body']['bottleneck_size']]
+config['cond_ae_body']['decoder_input'] = (4, 2, config['cond_ae_body']['bottleneck_size'])
+config['cond_ae_body']['decoder_depths'] = [3, 3, 9, 3]
+config['cond_ae_body']['decoder_dims'] = [config['cond_ae_body']['bottleneck_size'], 128, 64, 16]

configs/config_abl_2_pyr_32_64_128_256_512_skirt.py

@@ -0,0 +1,189 @@
+config = dict()
+
+config['gpu'] = '5'
+config['name'] = 'abl_2_pyr_32_64_128_256_512_skirt'
+config['notes'] = ''
+config['tags'] = ['pyr_skirt_2']
+
+# DON'T FORGET TO RESET SSIM LOSS USAGE
+# DON'T FORGET TO CHECK MODEL LOADING
+# DON'T FORGET TO CHECK IF STORE IN MEMORY
+# DON'T FORGET TO CHECK REG LOSS COEFF
+# DON'T FORGET TO LOAD FINAL UPSCALER UNLESS HAVE LEVEL 512
+# DON'T FORGET TO CHECK IF SKIRT
+# DON'T FORGET TO CHECK IF CAPE
+# DON'T FORGET TO CHECK IF USING DOWNSCALE NONE OR _2
+# DON'T FORGET TO CHECK UPSCALING and ENCODER CHECKPOINTS
+
+config['use_skinning'] =              True # and deshaping
+config['remove_global_orientation'] = True
+config['use_extra_losses'] =          True # edge, collision, 3d mesh
+config['use_offsets'] =               False
+
+config['use_norm_loss'] =             False
+
+config['pretrained_model'] = ''
+config['pretrained_model_iteration'] = ''
+
+config['pyramid'] = dict()
+config['pyramid']['resolutions'] = [32, 64, 128, 256, 512]
+config['pyramid']['pretrained_pyramid'] = 'checkpoints/abl_2_pyr_32_64_128_256_skirt'
+config['pyramid']['pretrained_pyramid_iteration'] = '70000'
+config['pyramid']['final_upscale_enhancer'] = ''
+
+config['train_dir'] = '/home/workspace/local-data/train_skirt'
+config['valid_dir'] = '/home/workspace/local-data/valid_skirt'
+config['test_dir'] = '/home/workspace/data/dataset/test_skirt_ready_2'
+config['batch_size'] = 4
+config['is_skirt'] = True
+config['is_cape'] = False
+
+config['pyramid']['static_model'] = 'checkpoints/static_1_skirt'
+config['pyramid']['body_model'] = 'checkpoints/body_1_skirt_2_2'
+config['pyramid']['curvature_model'] = 'checkpoints/curv_1_skirt'
+
+config['dataset_in_memory'] = False
+config['resolution'] = None
+config['latent_dim'] = 128
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['lr'] = 0.0001
+config['starting_balancing_coeff'] = 0.0001
+config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
+config['avg_garment_path'] = ''
+
+config['signature_dynamic'] = { 'id': (), 'frame_nr': (),
+                          'pose_shape': (72 + 11,),
+                          'uv_body_posed': (512, 256, 3),
+                          'mesh_body_posed': (config['smpl_res'], 3,),
+                          'mesh_posed': (config['smpl_res'], 3,),
+                          'mesh_unposed': (config['smpl_res'], 3,),
+                          'uv_unposed': (512, 256, 3),
+                          'uv_curvature': (512, 256, 7),
+                        }
+config['signature_static'] = { 'uv_static': (512, 256, 3),
+                              'mesh_mask': (config['smpl_res'],),
+                              'uv_mask': (512, 256,),
+                              }
+
+config['pyramid']['is_label'] = False
+config['pyramid']['latent_dim'] = config['latent_dim']
+config['pyramid']['add_noise'] = True
+config['pyramid']['noise_amount'] = 0.001
+
+config['pyramid']['loss_weights'] = [1,1,1,1,1]
+
+for res in config['pyramid']['resolutions']:
+  if res != 512:
+      res_modifier = '_' + str(res)
+      config['signature_dynamic']['uv_body_posed' + res_modifier] = (res, res//2, 3)
+      config['signature_dynamic']['uv_unposed' + res_modifier] = (res, res//2, 3)
+      config['signature_dynamic']['uv_curvature' + res_modifier] = (res, res//2, 7)
+      config['signature_static']['uv_static' + res_modifier] = (res, res//2, 3)
+      config['signature_static']['uv_mask' + res_modifier] = (res, res//2)
+
+
+
+config['pyramid']['freeze_encoders'] = True
+
+config['pyramid']['levels'] = []
+for res in config['pyramid']['resolutions']:
+  if res == 512:
+    level = dict()
+    level['resolution'] = 512
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 7
+    level['encoder_depths'] = [3, 3, 9, 3]
+    level['encoder_dims'] = [32, 64, 128, 256]
+    level['decoder_input'] = (16, 8, 80)
+    level['decoder_depths'] = [3, 3, 9, 3]
+    level['decoder_dims'] = [256, 128, 64, 16]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = ''
+    config['pyramid']['levels'].append(level)
+  elif res == 256:
+    level = dict()
+    level['resolution'] = 256
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 7
+    level['encoder_depths'] = [9, 9, 9, 9]
+    level['encoder_dims'] = [32, 64, 128, 256]
+    level['decoder_input'] = (8, 4, 128)
+    level['decoder_depths'] = [9, 9, 9, 9]
+    level['decoder_dims'] = [256, 128, 64, 16]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = 'checkpoints/256_512_skirt'#'checkpoints/up_3_d256_512_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  elif res == 128:
+    level = dict()
+    level['resolution'] = 128
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 5
+    level['encoder_depths'] = [16, 9, 9, 9]
+    level['encoder_dims'] = [64, 128, 128, 256]
+    level['decoder_input'] = (4, 2, 128)
+    level['decoder_depths'] = [16, 9, 9, 9]
+    level['decoder_dims'] = [256, 128, 128, 64]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = 'checkpoints/128_512_skirt'#'checkpoints/up_3_d128_256_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  elif res == 64:
+    level = dict()
+    level['resolution'] = 64
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 5
+    level['encoder_depths'] = [9, 9, 9]
+    level['encoder_dims'] = [64, 128, 256]
+    level['decoder_input'] = (4, 2, 256)
+    level['decoder_depths'] = [9, 9, 9]
+    level['decoder_dims'] = [256, 128, 64]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = 'checkpoints/64_512_skirt'#checkpoints/up_3_d64_128_2'
+    level['upscale_enhancer_freeze'] = False
+    config['pyramid']['levels'].append(level)
+  elif res == 32:
+    level = dict()
+    level['resolution'] = 32
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 3
+    level['encoder_depths'] = [9, 9]
+    level['encoder_dims'] = [256, 256]
+    level['decoder_input'] = (4, 2, 256)
+    level['decoder_depths'] = [9, 9]
+    level['decoder_dims'] = [256, 256]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = 'checkpoints/32_512_skirt'#'checkpoints/up_3_d32_64_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  else:
+    raise NotImplementedError
+
+
+
+config['cond_ae'] = dict()
+config['cond_ae']['is_label'] = False
+config['cond_ae']['upscale_output'] = False
+config['cond_ae']['latent_dim'] = config['latent_dim']
+config['cond_ae']['kernel_size'] = 7
+config['cond_ae']['bottleneck_size'] = 64
+config['cond_ae']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae']['encoder_dims'] = [32, 64, 128, 256, config['cond_ae']['bottleneck_size']]
+config['cond_ae']['decoder_input'] = (8, 4, config['cond_ae']['bottleneck_size'])
+config['cond_ae']['decoder_depths'] = [3, 3, 9, 3, 3]
+config['cond_ae']['decoder_dims'] = [config['cond_ae']['bottleneck_size'], 256, 128, 64, 16]
+
+config['cond_ae_body'] = dict()
+config['cond_ae_body']['latent_dim'] = config['latent_dim']
+config['cond_ae_body']['kernel_size'] = 7
+config['cond_ae_body']['upscale_output'] = True
+config['cond_ae_body']['bottleneck_size'] = 512
+config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3, 3]
+config['cond_ae_body']['encoder_dims'] = [32, 64, 128, 256, 256, config['cond_ae_body']['bottleneck_size']]
+config['cond_ae_body']['decoder_input'] = (4, 2, config['cond_ae_body']['bottleneck_size'])
+config['cond_ae_body']['decoder_depths'] = [3, 3, 9, 3]
+config['cond_ae_body']['decoder_dims'] = [config['cond_ae_body']['bottleneck_size'], 128, 64, 16]

configs/config_abl_2_pyr_32_64_128_256_512_test.py

@@ -0,0 +1,190 @@
+config = dict()
+
+config['gpu'] = '4'
+config['name'] = 'abl_2_pyr_32_64_128_256_512_test'
+config['notes'] = 'fixed noise handling'
+config['tags'] = ['pyr_non_skirt_2']
+
+# DON'T FORGET TO RESET SSIM LOSS USAGE
+# DON'T FORGET TO CHECK MODEL LOADING
+# DON'T FORGET TO CHECK IF STORE IN MEMORY
+# DON'T FORGET TO CHECK REG LOSS COEFF
+# DON'T FORGET TO LOAD FINAL UPSCALER UNLESS HAVE LEVEL 512
+# DON'T FORGET TO CHECK IF SKIRT
+# DON'T FORGET TO CHECK IF CAPE
+# DON'T FORGET TO CHECK IF USING DOWNSCALE NONE OR _2
+# DON'T FORGET TO CHECK UPSCALING and ENCODER CHECKPOINTS
+
+config['use_skinning'] =              True # and deshaping
+config['remove_global_orientation'] = True
+config['use_extra_losses'] =          True # edge, collision, 3d mesh
+config['use_offsets'] =               False
+config['use_norm_loss'] =             False
+
+config['pretrained_model'] = ''
+config['pretrained_model_iteration'] = ''
+
+config['pyramid'] = dict()
+config['pyramid']['resolutions'] = [256]
+config['pyramid']['pretrained_pyramid'] = ''
+config['pyramid']['pretrained_pyramid_iteration'] = ''
+config['pyramid']['final_upscale_enhancer'] = 'x'
+
+config['frames_file'] = 'frames_00761_0.txt'
+config['frames_file_valid'] = 'frames_00761_0.txt'
+config['frames_file_test'] = 'frames_00761_0.txt'
+
+config['train_dir'] = '/home/workspace/local-data/train'
+config['valid_dir'] = '/home/workspace/local-data/train'
+config['test_dir'] = '/home/workspace/data/dataset/test_ready'
+config['batch_size'] = 4
+config['is_skirt'] = False
+config['is_cape'] = False
+
+config['dataset_in_memory'] = True
+config['resolution'] = None
+config['latent_dim'] = 128
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['lr'] = 0.0001
+config['starting_balancing_coeff'] = 0.0001
+config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
+config['avg_garment_path'] = ''
+
+config['signature_dynamic'] = { 'id': (), 'frame_nr': (),
+                          'pose_shape': (72 + 11,),
+                          'uv_body_posed': (512, 256, 3),
+                          'mesh_body_posed': (config['smpl_res'], 3,),
+                          'mesh_posed': (config['smpl_res'], 3,),
+                          'mesh_unposed': (config['smpl_res'], 3,),
+                          'uv_unposed': (512, 256, 3),
+                          'uv_curvature': (512, 256, 7),
+                        }
+config['signature_static'] = { 'uv_static': (512, 256, 3),
+                              'mesh_mask': (config['smpl_res'],),
+                              'uv_mask': (512, 256,),
+                              }
+
+config['pyramid']['is_label'] = False
+config['pyramid']['latent_dim'] = config['latent_dim']
+config['pyramid']['add_noise'] = True
+config['pyramid']['noise_amount'] = 0.001
+
+config['pyramid']['loss_weights'] = [1,1,1,1,1]
+
+for res in config['pyramid']['resolutions']:
+  if res != 512:
+      res_modifier = '_' + str(res)
+      config['signature_dynamic']['uv_body_posed' + res_modifier] = (res, res//2, 3)
+      config['signature_dynamic']['uv_unposed' + res_modifier] = (res, res//2, 3)
+      config['signature_dynamic']['uv_curvature' + res_modifier] = (res, res//2, 7)
+      config['signature_static']['uv_static' + res_modifier] = (res, res//2, 3)
+      config['signature_static']['uv_mask' + res_modifier] = (res, res//2)
+
+
+config['pyramid']['static_model'] = 'checkpoints/static_1'
+config['pyramid']['body_model'] = 'checkpoints/body_1'
+config['pyramid']['curvature_model'] = 'checkpoints/curv_1'
+config['pyramid']['freeze_encoders'] = True
+
+config['pyramid']['levels'] = []
+for res in config['pyramid']['resolutions']:
+  if res == 512:
+    level = dict()
+    level['resolution'] = 512
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 7
+    level['encoder_depths'] = [3, 3, 9, 3]
+    level['encoder_dims'] = [32, 64, 128, 256]
+    level['decoder_input'] = (16, 8, 80)
+    level['decoder_depths'] = [3, 3, 9, 3]
+    level['decoder_dims'] = [256, 128, 64, 16]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = ''
+    config['pyramid']['levels'].append(level)
+  elif res == 256:
+    level = dict()
+    level['resolution'] = 256
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 7
+    level['encoder_depths'] = [9, 9, 9, 9]
+    level['encoder_dims'] = [32, 64, 128, 256]
+    level['decoder_input'] = (8, 4, 128)
+    level['decoder_depths'] = [9, 9, 9, 9]
+    level['decoder_dims'] = [256, 128, 64, 16]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = ''#'checkpoints/up_3_d256_512_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  elif res == 128:
+    level = dict()
+    level['resolution'] = 128
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 5
+    level['encoder_depths'] = [16, 9, 9, 9]
+    level['encoder_dims'] = [64, 128, 128, 256]
+    level['decoder_input'] = (4, 2, 128)
+    level['decoder_depths'] = [16, 9, 9, 9]
+    level['decoder_dims'] = [256, 128, 128, 64]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = ''#'checkpoints/up_3_d128_256_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  elif res == 64:
+    level = dict()
+    level['resolution'] = 64
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 5
+    level['encoder_depths'] = [9, 9, 9]
+    level['encoder_dims'] = [64, 128, 256]
+    level['decoder_input'] = (4, 2, 256)
+    level['decoder_depths'] = [9, 9, 9]
+    level['decoder_dims'] = [256, 128, 64]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = ''#checkpoints/up_3_d64_128_2'
+    level['upscale_enhancer_freeze'] = False
+    config['pyramid']['levels'].append(level)
+  elif res == 32:
+    level = dict()
+    level['resolution'] = 32
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 3
+    level['encoder_depths'] = [9, 9]
+    level['encoder_dims'] = [256, 256]
+    level['decoder_input'] = (4, 2, 256)
+    level['decoder_depths'] = [9, 9]
+    level['decoder_dims'] = [256, 256]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = ''#'checkpoints/up_3_d32_64_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  else:
+    raise NotImplementedError
+
+
+
+config['cond_ae'] = dict()
+config['cond_ae']['is_label'] = False
+config['cond_ae']['upscale_output'] = False
+config['cond_ae']['latent_dim'] = config['latent_dim']
+config['cond_ae']['kernel_size'] = 7
+config['cond_ae']['bottleneck_size'] = 64
+config['cond_ae']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae']['encoder_dims'] = [32, 64, 128, 256, config['cond_ae']['bottleneck_size']]
+config['cond_ae']['decoder_input'] = (8, 4, config['cond_ae']['bottleneck_size'])
+config['cond_ae']['decoder_depths'] = [3, 3, 9, 3, 3]
+config['cond_ae']['decoder_dims'] = [config['cond_ae']['bottleneck_size'], 256, 128, 64, 16]
+
+config['cond_ae_body'] = dict()
+config['cond_ae_body']['latent_dim'] = config['latent_dim']
+config['cond_ae_body']['kernel_size'] = 7
+config['cond_ae_body']['upscale_output'] = True
+config['cond_ae_body']['bottleneck_size'] = 512
+config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3, 3]
+config['cond_ae_body']['encoder_dims'] = [32, 64, 128, 256, 256, config['cond_ae_body']['bottleneck_size']]
+config['cond_ae_body']['decoder_input'] = (4, 2, config['cond_ae_body']['bottleneck_size'])
+config['cond_ae_body']['decoder_depths'] = [3, 3, 9, 3]
+config['cond_ae_body']['decoder_dims'] = [config['cond_ae_body']['bottleneck_size'], 128, 64, 16]

configs/config_abl_2_pyr_32_64_128_256_skirt.py

@@ -0,0 +1,189 @@
+config = dict()
+
+config['gpu'] = '4'
+config['name'] = 'abl_2_pyr_32_64_128_256_skirt'
+config['notes'] = ''
+config['tags'] = ['pyr_skirt_2']
+
+# DON'T FORGET TO RESET SSIM LOSS USAGE
+# DON'T FORGET TO CHECK MODEL LOADING
+# DON'T FORGET TO CHECK IF STORE IN MEMORY
+# DON'T FORGET TO CHECK REG LOSS COEFF
+# DON'T FORGET TO LOAD FINAL UPSCALER UNLESS HAVE LEVEL 512
+# DON'T FORGET TO CHECK IF SKIRT
+# DON'T FORGET TO CHECK IF CAPE
+# DON'T FORGET TO CHECK IF USING DOWNSCALE NONE OR _2
+# DON'T FORGET TO CHECK UPSCALING and ENCODER CHECKPOINTS
+
+config['use_skinning'] =              True # and deshaping
+config['remove_global_orientation'] = True
+config['use_extra_losses'] =          True # edge, collision, 3d mesh
+config['use_offsets'] =               False
+
+config['use_norm_loss'] =             False
+
+config['pretrained_model'] = ''
+config['pretrained_model_iteration'] = ''
+
+config['pyramid'] = dict()
+config['pyramid']['resolutions'] = [32, 64, 128, 256]
+config['pyramid']['pretrained_pyramid'] = 'checkpoints/abl_2_pyr_32_64_128_skirt'
+config['pyramid']['pretrained_pyramid_iteration'] = '80000'
+config['pyramid']['final_upscale_enhancer'] = 'checkpoints/256_512_skirt'
+
+config['train_dir'] = '/home/workspace/local-data/train_skirt'
+config['valid_dir'] = '/home/workspace/local-data/valid_skirt'
+config['test_dir'] = '/home/workspace/data/dataset/test_skirt_ready_2'
+config['batch_size'] = 4
+config['is_skirt'] = True
+config['is_cape'] = False
+
+config['pyramid']['static_model'] = 'checkpoints/static_1_skirt'
+config['pyramid']['body_model'] = 'checkpoints/body_1_skirt_2_2'
+config['pyramid']['curvature_model'] = 'checkpoints/curv_1_skirt'
+
+config['dataset_in_memory'] = False
+config['resolution'] = None
+config['latent_dim'] = 128
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['lr'] = 0.0001
+config['starting_balancing_coeff'] = 0.0001
+config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
+config['avg_garment_path'] = ''
+
+config['signature_dynamic'] = { 'id': (), 'frame_nr': (),
+                          'pose_shape': (72 + 11,),
+                          'uv_body_posed': (512, 256, 3),
+                          'mesh_body_posed': (config['smpl_res'], 3,),
+                          'mesh_posed': (config['smpl_res'], 3,),
+                          'mesh_unposed': (config['smpl_res'], 3,),
+                          'uv_unposed': (512, 256, 3),
+                          'uv_curvature': (512, 256, 7),
+                        }
+config['signature_static'] = { 'uv_static': (512, 256, 3),
+                              'mesh_mask': (config['smpl_res'],),
+                              'uv_mask': (512, 256,),
+                              }
+
+config['pyramid']['is_label'] = False
+config['pyramid']['latent_dim'] = config['latent_dim']
+config['pyramid']['add_noise'] = True
+config['pyramid']['noise_amount'] = 0.001
+
+config['pyramid']['loss_weights'] = [1,1,1,1,1]
+
+for res in config['pyramid']['resolutions']:
+  if res != 512:
+      res_modifier = '_' + str(res)
+      config['signature_dynamic']['uv_body_posed' + res_modifier] = (res, res//2, 3)
+      config['signature_dynamic']['uv_unposed' + res_modifier] = (res, res//2, 3)
+      config['signature_dynamic']['uv_curvature' + res_modifier] = (res, res//2, 7)
+      config['signature_static']['uv_static' + res_modifier] = (res, res//2, 3)
+      config['signature_static']['uv_mask' + res_modifier] = (res, res//2)
+
+
+
+config['pyramid']['freeze_encoders'] = True
+
+config['pyramid']['levels'] = []
+for res in config['pyramid']['resolutions']:
+  if res == 512:
+    level = dict()
+    level['resolution'] = 512
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 7
+    level['encoder_depths'] = [3, 3, 9, 3]
+    level['encoder_dims'] = [32, 64, 128, 256]
+    level['decoder_input'] = (16, 8, 80)
+    level['decoder_depths'] = [3, 3, 9, 3]
+    level['decoder_dims'] = [256, 128, 64, 16]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = ''
+    config['pyramid']['levels'].append(level)
+  elif res == 256:
+    level = dict()
+    level['resolution'] = 256
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 7
+    level['encoder_depths'] = [9, 9, 9, 9]
+    level['encoder_dims'] = [32, 64, 128, 256]
+    level['decoder_input'] = (8, 4, 128)
+    level['decoder_depths'] = [9, 9, 9, 9]
+    level['decoder_dims'] = [256, 128, 64, 16]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = ''#'checkpoints/up_3_d256_512_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  elif res == 128:
+    level = dict()
+    level['resolution'] = 128
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 5
+    level['encoder_depths'] = [16, 9, 9, 9]
+    level['encoder_dims'] = [64, 128, 128, 256]
+    level['decoder_input'] = (4, 2, 128)
+    level['decoder_depths'] = [16, 9, 9, 9]
+    level['decoder_dims'] = [256, 128, 128, 64]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = 'checkpoints/128_512_skirt'#'checkpoints/up_3_d128_256_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  elif res == 64:
+    level = dict()
+    level['resolution'] = 64
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 5
+    level['encoder_depths'] = [9, 9, 9]
+    level['encoder_dims'] = [64, 128, 256]
+    level['decoder_input'] = (4, 2, 256)
+    level['decoder_depths'] = [9, 9, 9]
+    level['decoder_dims'] = [256, 128, 64]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = 'checkpoints/64_512_skirt'#checkpoints/up_3_d64_128_2'
+    level['upscale_enhancer_freeze'] = False
+    config['pyramid']['levels'].append(level)
+  elif res == 32:
+    level = dict()
+    level['resolution'] = 32
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 3
+    level['encoder_depths'] = [9, 9]
+    level['encoder_dims'] = [256, 256]
+    level['decoder_input'] = (4, 2, 256)
+    level['decoder_depths'] = [9, 9]
+    level['decoder_dims'] = [256, 256]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = 'checkpoints/32_512_skirt'#'checkpoints/up_3_d32_64_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  else:
+    raise NotImplementedError
+
+
+
+config['cond_ae'] = dict()
+config['cond_ae']['is_label'] = False
+config['cond_ae']['upscale_output'] = False
+config['cond_ae']['latent_dim'] = config['latent_dim']
+config['cond_ae']['kernel_size'] = 7
+config['cond_ae']['bottleneck_size'] = 64
+config['cond_ae']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae']['encoder_dims'] = [32, 64, 128, 256, config['cond_ae']['bottleneck_size']]
+config['cond_ae']['decoder_input'] = (8, 4, config['cond_ae']['bottleneck_size'])
+config['cond_ae']['decoder_depths'] = [3, 3, 9, 3, 3]
+config['cond_ae']['decoder_dims'] = [config['cond_ae']['bottleneck_size'], 256, 128, 64, 16]
+
+config['cond_ae_body'] = dict()
+config['cond_ae_body']['latent_dim'] = config['latent_dim']
+config['cond_ae_body']['kernel_size'] = 7
+config['cond_ae_body']['upscale_output'] = True
+config['cond_ae_body']['bottleneck_size'] = 512
+config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3, 3]
+config['cond_ae_body']['encoder_dims'] = [32, 64, 128, 256, 256, config['cond_ae_body']['bottleneck_size']]
+config['cond_ae_body']['decoder_input'] = (4, 2, config['cond_ae_body']['bottleneck_size'])
+config['cond_ae_body']['decoder_depths'] = [3, 3, 9, 3]
+config['cond_ae_body']['decoder_dims'] = [config['cond_ae_body']['bottleneck_size'], 128, 64, 16]

configs/config_abl_2_pyr_32_64_128_skirt.py

@@ -0,0 +1,189 @@
+config = dict()
+
+config['gpu'] = '5'
+config['name'] = 'abl_2_pyr_32_64_128_skirt'
+config['notes'] = ''
+config['tags'] = ['pyr_skirt_2']
+
+# DON'T FORGET TO RESET SSIM LOSS USAGE
+# DON'T FORGET TO CHECK MODEL LOADING
+# DON'T FORGET TO CHECK IF STORE IN MEMORY
+# DON'T FORGET TO CHECK REG LOSS COEFF
+# DON'T FORGET TO LOAD FINAL UPSCALER UNLESS HAVE LEVEL 512
+# DON'T FORGET TO CHECK IF SKIRT
+# DON'T FORGET TO CHECK IF CAPE
+# DON'T FORGET TO CHECK IF USING DOWNSCALE NONE OR _2
+# DON'T FORGET TO CHECK UPSCALING and ENCODER CHECKPOINTS
+
+config['use_skinning'] =              True # and deshaping
+config['remove_global_orientation'] = True
+config['use_extra_losses'] =          True # edge, collision, 3d mesh
+config['use_offsets'] =               False
+
+config['use_norm_loss'] =             False
+
+config['pretrained_model'] = ''
+config['pretrained_model_iteration'] = ''
+
+config['pyramid'] = dict()
+config['pyramid']['resolutions'] = [32, 64, 128]
+config['pyramid']['pretrained_pyramid'] = 'checkpoints/abl_2_pyr_32_64_skirt'
+config['pyramid']['pretrained_pyramid_iteration'] = ''
+config['pyramid']['final_upscale_enhancer'] = 'checkpoints/128_512_skirt'
+
+config['train_dir'] = '/home/workspace/local-data/train_skirt'
+config['valid_dir'] = '/home/workspace/local-data/valid_skirt'
+config['test_dir'] = '/home/workspace/data/dataset/test_skirt_ready_2'
+config['batch_size'] = 4
+config['is_skirt'] = True
+config['is_cape'] = False
+
+config['pyramid']['static_model'] = 'checkpoints/static_1_skirt'
+config['pyramid']['body_model'] = 'checkpoints/body_1_skirt_2_2'
+config['pyramid']['curvature_model'] = 'checkpoints/curv_1_skirt'
+
+config['dataset_in_memory'] = False
+config['resolution'] = None
+config['latent_dim'] = 128
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['lr'] = 0.0001
+config['starting_balancing_coeff'] = 0.0001
+config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
+config['avg_garment_path'] = ''
+
+config['signature_dynamic'] = { 'id': (), 'frame_nr': (),
+                          'pose_shape': (72 + 11,),
+                          'uv_body_posed': (512, 256, 3),
+                          'mesh_body_posed': (config['smpl_res'], 3,),
+                          'mesh_posed': (config['smpl_res'], 3,),
+                          'mesh_unposed': (config['smpl_res'], 3,),
+                          'uv_unposed': (512, 256, 3),
+                          'uv_curvature': (512, 256, 7),
+                        }
+config['signature_static'] = { 'uv_static': (512, 256, 3),
+                              'mesh_mask': (config['smpl_res'],),
+                              'uv_mask': (512, 256,),
+                              }
+
+config['pyramid']['is_label'] = False
+config['pyramid']['latent_dim'] = config['latent_dim']
+config['pyramid']['add_noise'] = True
+config['pyramid']['noise_amount'] = 0.001
+
+config['pyramid']['loss_weights'] = [1,1,1,1,1]
+
+for res in config['pyramid']['resolutions']:
+  if res != 512:
+      res_modifier = '_' + str(res)
+      config['signature_dynamic']['uv_body_posed' + res_modifier] = (res, res//2, 3)
+      config['signature_dynamic']['uv_unposed' + res_modifier] = (res, res//2, 3)
+      config['signature_dynamic']['uv_curvature' + res_modifier] = (res, res//2, 7)
+      config['signature_static']['uv_static' + res_modifier] = (res, res//2, 3)
+      config['signature_static']['uv_mask' + res_modifier] = (res, res//2)
+
+
+
+config['pyramid']['freeze_encoders'] = True
+
+config['pyramid']['levels'] = []
+for res in config['pyramid']['resolutions']:
+  if res == 512:
+    level = dict()
+    level['resolution'] = 512
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 7
+    level['encoder_depths'] = [3, 3, 9, 3]
+    level['encoder_dims'] = [32, 64, 128, 256]
+    level['decoder_input'] = (16, 8, 80)
+    level['decoder_depths'] = [3, 3, 9, 3]
+    level['decoder_dims'] = [256, 128, 64, 16]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = ''
+    config['pyramid']['levels'].append(level)
+  elif res == 256:
+    level = dict()
+    level['resolution'] = 256
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 7
+    level['encoder_depths'] = [9, 9, 9, 9]
+    level['encoder_dims'] = [32, 64, 128, 256]
+    level['decoder_input'] = (8, 4, 128)
+    level['decoder_depths'] = [9, 9, 9, 9]
+    level['decoder_dims'] = [256, 128, 64, 16]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = ''#'checkpoints/up_3_d256_512_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  elif res == 128:
+    level = dict()
+    level['resolution'] = 128
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 5
+    level['encoder_depths'] = [16, 9, 9, 9]
+    level['encoder_dims'] = [64, 128, 128, 256]
+    level['decoder_input'] = (4, 2, 128)
+    level['decoder_depths'] = [16, 9, 9, 9]
+    level['decoder_dims'] = [256, 128, 128, 64]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = ''#'checkpoints/up_3_d128_256_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  elif res == 64:
+    level = dict()
+    level['resolution'] = 64
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 5
+    level['encoder_depths'] = [9, 9, 9]
+    level['encoder_dims'] = [64, 128, 256]
+    level['decoder_input'] = (4, 2, 256)
+    level['decoder_depths'] = [9, 9, 9]
+    level['decoder_dims'] = [256, 128, 64]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = 'checkpoints/64_512_skirt'#checkpoints/up_3_d64_128_2'
+    level['upscale_enhancer_freeze'] = False
+    config['pyramid']['levels'].append(level)
+  elif res == 32:
+    level = dict()
+    level['resolution'] = 32
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 3
+    level['encoder_depths'] = [9, 9]
+    level['encoder_dims'] = [256, 256]
+    level['decoder_input'] = (4, 2, 256)
+    level['decoder_depths'] = [9, 9]
+    level['decoder_dims'] = [256, 256]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = 'checkpoints/32_512_skirt'#'checkpoints/up_3_d32_64_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  else:
+    raise NotImplementedError
+
+
+
+config['cond_ae'] = dict()
+config['cond_ae']['is_label'] = False
+config['cond_ae']['upscale_output'] = False
+config['cond_ae']['latent_dim'] = config['latent_dim']
+config['cond_ae']['kernel_size'] = 7
+config['cond_ae']['bottleneck_size'] = 64
+config['cond_ae']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae']['encoder_dims'] = [32, 64, 128, 256, config['cond_ae']['bottleneck_size']]
+config['cond_ae']['decoder_input'] = (8, 4, config['cond_ae']['bottleneck_size'])
+config['cond_ae']['decoder_depths'] = [3, 3, 9, 3, 3]
+config['cond_ae']['decoder_dims'] = [config['cond_ae']['bottleneck_size'], 256, 128, 64, 16]
+
+config['cond_ae_body'] = dict()
+config['cond_ae_body']['latent_dim'] = config['latent_dim']
+config['cond_ae_body']['kernel_size'] = 7
+config['cond_ae_body']['upscale_output'] = True
+config['cond_ae_body']['bottleneck_size'] = 512
+config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3, 3]
+config['cond_ae_body']['encoder_dims'] = [32, 64, 128, 256, 256, config['cond_ae_body']['bottleneck_size']]
+config['cond_ae_body']['decoder_input'] = (4, 2, config['cond_ae_body']['bottleneck_size'])
+config['cond_ae_body']['decoder_depths'] = [3, 3, 9, 3]
+config['cond_ae_body']['decoder_dims'] = [config['cond_ae_body']['bottleneck_size'], 128, 64, 16]

configs/config_abl_2_pyr_32_64_skirt.py

@@ -1,6 +1,6 @@
 config = dict()
 
-config['gpu'] = '4'
+config['gpu'] = '6'
 config['name'] = 'abl_2_pyr_32_64_skirt'
 config['notes'] = ''
 config['tags'] = ['pyr_skirt_2']

configs/config_abl_2_pyr_32_skirt.py

@@ -1,6 +1,6 @@
 config = dict()
 
-config['gpu'] = '4'
+config['gpu'] = '6'
 config['name'] = 'abl_2_pyr_32_skirt'
 config['notes'] = ''
 config['tags'] = ['pyr_skirt_2']

configs/config_abl_2_pyr_64_128_256_512_cape_f.py

@@ -1,6 +1,6 @@
 config = dict()
 
-config['gpu'] = '6'
+config['gpu'] = '2'
 config['name'] = 'abl_2_pyr_64_128_256_512_cape_f'
 config['notes'] = ''
 config['tags'] = ['cape_f_2']
@@ -116,6 +116,7 @@ for res in config['pyramid']['resolutions']:
     level['avg_garment_path'] = ''
     level['upscale_enhancer'] = 'checkpoints/256_512_cape_f'#'checkpoints/up_3_d256_512_2'
     level['upscale_enhancer_freeze'] = True
+    level['upscale_enhancer_kernel_size'] = 3
     config['pyramid']['levels'].append(level)
   elif res == 128:
     level = dict()
@@ -130,6 +131,7 @@ for res in config['pyramid']['resolutions']:
     level['avg_garment_path'] = ''
     level['upscale_enhancer'] = 'checkpoints/128_512_cape_f'#'checkpoints/up_3_d128_256_2'
     level['upscale_enhancer_freeze'] = True
+    level['upscale_enhancer_kernel_size'] = 3
     config['pyramid']['levels'].append(level)
   elif res == 64:
     level = dict()
@@ -144,6 +146,7 @@ for res in config['pyramid']['resolutions']:
     level['avg_garment_path'] = ''
     level['upscale_enhancer'] = 'checkpoints/64_512_cape_f'#checkpoints/up_3_d64_128_2'
     level['upscale_enhancer_freeze'] = False
+    level['upscale_enhancer_kernel_size'] = 3
     config['pyramid']['levels'].append(level)
   elif res == 32:
     level = dict()

configs/config_abl_2_pyr_64_128_256_cape_m_00032_SL.py

@@ -0,0 +1,190 @@
+config = dict()
+
+config['gpu'] = '4'
+config['name'] = 'abl_2_pyr_64_128_256_cape_m_00032_SL'
+config['notes'] = ''
+config['tags'] = ['cape_m_2']
+
+# DON'T FORGET TO RESET SSIM LOSS USAGE
+# DON'T FORGET TO CHECK MODEL LOADING
+# DON'T FORGET TO CHECK IF STORE IN MEMORY
+# DON'T FORGET TO CHECK REG LOSS COEFF
+# DON'T FORGET TO LOAD FINAL UPSCALER UNLESS HAVE LEVEL 512
+# DON'T FORGET TO CHECK IF SKIRT
+# DON'T FORGET TO CHECK IF CAPE
+# DON'T FORGET TO CHECK IF USING DOWNSCALE NONE OR _2
+# DON'T FORGET TO CHECK UPSCALING and ENCODER CHECKPOINTS
+
+config['use_skinning'] =              True # and deshaping
+config['remove_global_orientation'] = True
+config['use_extra_losses'] =          True # edge, collision, 3d mesh
+config['use_offsets'] =               False
+config['use_norm_loss'] =             False
+
+config['pretrained_model'] = ''
+config['pretrained_model_iteration'] = ''
+
+config['pyramid'] = dict()
+config['pyramid']['resolutions'] = [64, 128, 256]
+config['pyramid']['pretrained_pyramid'] = 'checkpoints/abl_2_pyr_64_128_cape_m_00032_SL'
+config['pyramid']['final_upscale_enhancer'] = 'checkpoints/256_512_cape_m'
+
+config['train_dir'] = '/home/workspace/local-data/cape_processed/train/male'
+config['valid_dir'] = '/home/workspace/local-data/cape_processed/train/male'
+config['test_dir'] = '/home/workspace/local-data/cape_processed/train/male'
+config['batch_size'] = 4
+config['is_skirt'] = False
+config['is_cape'] = True
+
+config['frames_file'] = 'frames_00032_SL_train.txt'
+config['frames_file_valid'] = 'frames_00032_SL_test.txt'
+
+config['pyramid']['static_model'] = 'checkpoints/static_1_cape_m'
+config['pyramid']['body_model'] = 'checkpoints/body_1_cape_m'
+config['pyramid']['curvature_model'] = 'checkpoints/curv_1_cape_m'
+
+config['dataset_in_memory'] = True
+config['resolution'] = None
+config['latent_dim'] = 128
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['lr'] = 0.0001
+config['starting_balancing_coeff'] = 0.0001
+config['reg_loss_coeff'] = None # set to 0 if don't want to use edge+collision
+config['avg_garment_path'] = ''
+
+config['signature_dynamic'] = { 'id': (), 'frame_nr': (),
+                          'pose_shape': (72 + 11,),
+                          'uv_body_posed': (512, 256, 3),
+                          'mesh_body_posed': (config['smpl_res'], 3,),
+                          'mesh_posed': (config['smpl_res'], 3,),
+                          'mesh_unposed': (config['smpl_res'], 3,),
+                          'uv_unposed': (512, 256, 3),
+                          'uv_curvature': (512, 256, 7),
+                        }
+config['signature_static'] = { 'uv_static': (512, 256, 3),
+                              'mesh_mask': (config['smpl_res'],),
+                              'uv_mask': (512, 256,),
+                              }
+
+config['pyramid']['is_label'] = False
+config['pyramid']['latent_dim'] = config['latent_dim']
+config['pyramid']['add_noise'] = True
+config['pyramid']['noise_amount'] = 0.001
+
+config['pyramid']['loss_weights'] = [1,1,1,1,1]
+
+for res in config['pyramid']['resolutions']:
+  if res != 512:
+      res_modifier = '_' + str(res)
+      config['signature_dynamic']['uv_body_posed' + res_modifier] = (res, res//2, 3)
+      config['signature_dynamic']['uv_unposed' + res_modifier] = (res, res//2, 3)
+      config['signature_dynamic']['uv_curvature' + res_modifier] = (res, res//2, 7)
+      config['signature_static']['uv_static' + res_modifier] = (res, res//2, 3)
+      config['signature_static']['uv_mask' + res_modifier] = (res, res//2)
+
+
+
+config['pyramid']['freeze_encoders'] = True
+
+config['pyramid']['levels'] = []
+for res in config['pyramid']['resolutions']:
+  if res == 512:
+    level = dict()
+    level['resolution'] = 512
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 7
+    level['encoder_depths'] = [3, 3, 9, 3]
+    level['encoder_dims'] = [32, 64, 128, 256]
+    level['decoder_input'] = (16, 8, 80)
+    level['decoder_depths'] = [3, 3, 9, 3]
+    level['decoder_dims'] = [256, 128, 64, 16]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = ''
+    config['pyramid']['levels'].append(level)
+  elif res == 256:
+    level = dict()
+    level['resolution'] = 256
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 7
+    level['encoder_depths'] = [9, 9, 9, 9]
+    level['encoder_dims'] = [32, 64, 128, 256]
+    level['decoder_input'] = (8, 4, 128)
+    level['decoder_depths'] = [9, 9, 9, 9]
+    level['decoder_dims'] = [256, 128, 64, 16]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = ''#'checkpoints/up_3_d256_512_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  elif res == 128:
+    level = dict()
+    level['resolution'] = 128
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 5
+    level['encoder_depths'] = [16, 9, 9, 9]
+    level['encoder_dims'] = [64, 128, 128, 256]
+    level['decoder_input'] = (4, 2, 128)
+    level['decoder_depths'] = [16, 9, 9, 9]
+    level['decoder_dims'] = [256, 128, 128, 64]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = 'checkpoints/128_512_cape_m'#'checkpoints/up_3_d128_256_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  elif res == 64:
+    level = dict()
+    level['resolution'] = 64
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 5
+    level['encoder_depths'] = [9, 9, 9]
+    level['encoder_dims'] = [64, 128, 256]
+    level['decoder_input'] = (4, 2, 256)
+    level['decoder_depths'] = [9, 9, 9]
+    level['decoder_dims'] = [256, 128, 64]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = 'checkpoints/64_512_cape_m'#checkpoints/up_3_d64_128_2'
+    level['upscale_enhancer_freeze'] = False
+    config['pyramid']['levels'].append(level)
+  elif res == 32:
+    level = dict()
+    level['resolution'] = 32
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 3
+    level['encoder_depths'] = [9, 9]
+    level['encoder_dims'] = [256, 256]
+    level['decoder_input'] = (4, 2, 256)
+    level['decoder_depths'] = [9, 9]
+    level['decoder_dims'] = [256, 256]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = ''#'checkpoints/up_3_d32_64_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  else:
+    raise NotImplementedError
+
+
+
+config['cond_ae'] = dict()
+config['cond_ae']['is_label'] = False
+config['cond_ae']['upscale_output'] = False
+config['cond_ae']['latent_dim'] = config['latent_dim']
+config['cond_ae']['kernel_size'] = 7
+config['cond_ae']['bottleneck_size'] = 64
+config['cond_ae']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae']['encoder_dims'] = [32, 64, 128, 256, config['cond_ae']['bottleneck_size']]
+config['cond_ae']['decoder_input'] = (8, 4, config['cond_ae']['bottleneck_size'])
+config['cond_ae']['decoder_depths'] = [3, 3, 9, 3, 3]
+config['cond_ae']['decoder_dims'] = [config['cond_ae']['bottleneck_size'], 256, 128, 64, 16]
+
+config['cond_ae_body'] = dict()
+config['cond_ae_body']['latent_dim'] = config['latent_dim']
+config['cond_ae_body']['kernel_size'] = 7
+config['cond_ae_body']['upscale_output'] = True
+config['cond_ae_body']['bottleneck_size'] = 512
+config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3, 3]
+config['cond_ae_body']['encoder_dims'] = [32, 64, 128, 256, 256, config['cond_ae_body']['bottleneck_size']]
+config['cond_ae_body']['decoder_input'] = (4, 2, config['cond_ae_body']['bottleneck_size'])
+config['cond_ae_body']['decoder_depths'] = [3, 3, 9, 3]
+config['cond_ae_body']['decoder_dims'] = [config['cond_ae_body']['bottleneck_size'], 128, 64, 16]

configs/config_abl_2_pyr_64_128_cape_m_00032_SL.py

@@ -0,0 +1,190 @@
+config = dict()
+
+config['gpu'] = '2'
+config['name'] = 'abl_2_pyr_64_128_cape_m_00032_SL'
+config['notes'] = ''
+config['tags'] = ['cape_m_2']
+
+# DON'T FORGET TO RESET SSIM LOSS USAGE
+# DON'T FORGET TO CHECK MODEL LOADING
+# DON'T FORGET TO CHECK IF STORE IN MEMORY
+# DON'T FORGET TO CHECK REG LOSS COEFF
+# DON'T FORGET TO LOAD FINAL UPSCALER UNLESS HAVE LEVEL 512
+# DON'T FORGET TO CHECK IF SKIRT
+# DON'T FORGET TO CHECK IF CAPE
+# DON'T FORGET TO CHECK IF USING DOWNSCALE NONE OR _2
+# DON'T FORGET TO CHECK UPSCALING and ENCODER CHECKPOINTS
+
+config['use_skinning'] =              True # and deshaping
+config['remove_global_orientation'] = True
+config['use_extra_losses'] =          True # edge, collision, 3d mesh
+config['use_offsets'] =               False
+config['use_norm_loss'] =             False
+
+config['pretrained_model'] = ''
+config['pretrained_model_iteration'] = ''
+
+config['pyramid'] = dict()
+config['pyramid']['resolutions'] = [64, 128]
+config['pyramid']['pretrained_pyramid'] = 'checkpoints/abl_2_pyr_64_cape_m_00032_SL'
+config['pyramid']['final_upscale_enhancer'] = 'checkpoints/128_512_cape_m'
+
+config['train_dir'] = '/home/workspace/local-data/cape_processed/train/male'
+config['valid_dir'] = '/home/workspace/local-data/cape_processed/train/male'
+config['test_dir'] = '/home/workspace/local-data/cape_processed/train/male'
+config['batch_size'] = 4
+config['is_skirt'] = False
+config['is_cape'] = True
+
+config['frames_file'] = 'frames_00032_SL_train.txt'
+config['frames_file_valid'] = 'frames_00032_SL_test.txt'
+
+config['pyramid']['static_model'] = 'checkpoints/static_1_cape_m'
+config['pyramid']['body_model'] = 'checkpoints/body_1_cape_m'
+config['pyramid']['curvature_model'] = 'checkpoints/curv_1_cape_m'
+
+config['dataset_in_memory'] = True
+config['resolution'] = None
+config['latent_dim'] = 128
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['lr'] = 0.0001
+config['starting_balancing_coeff'] = 0.0001
+config['reg_loss_coeff'] = None # set to 0 if don't want to use edge+collision
+config['avg_garment_path'] = ''
+
+config['signature_dynamic'] = { 'id': (), 'frame_nr': (),
+                          'pose_shape': (72 + 11,),
+                          'uv_body_posed': (512, 256, 3),
+                          'mesh_body_posed': (config['smpl_res'], 3,),
+                          'mesh_posed': (config['smpl_res'], 3,),
+                          'mesh_unposed': (config['smpl_res'], 3,),
+                          'uv_unposed': (512, 256, 3),
+                          'uv_curvature': (512, 256, 7),
+                        }
+config['signature_static'] = { 'uv_static': (512, 256, 3),
+                              'mesh_mask': (config['smpl_res'],),
+                              'uv_mask': (512, 256,),
+                              }
+
+config['pyramid']['is_label'] = False
+config['pyramid']['latent_dim'] = config['latent_dim']
+config['pyramid']['add_noise'] = True
+config['pyramid']['noise_amount'] = 0.001
+
+config['pyramid']['loss_weights'] = [1,1,1,1,1]
+
+for res in config['pyramid']['resolutions']:
+  if res != 512:
+      res_modifier = '_' + str(res)
+      config['signature_dynamic']['uv_body_posed' + res_modifier] = (res, res//2, 3)
+      config['signature_dynamic']['uv_unposed' + res_modifier] = (res, res//2, 3)
+      config['signature_dynamic']['uv_curvature' + res_modifier] = (res, res//2, 7)
+      config['signature_static']['uv_static' + res_modifier] = (res, res//2, 3)
+      config['signature_static']['uv_mask' + res_modifier] = (res, res//2)
+
+
+
+config['pyramid']['freeze_encoders'] = True
+
+config['pyramid']['levels'] = []
+for res in config['pyramid']['resolutions']:
+  if res == 512:
+    level = dict()
+    level['resolution'] = 512
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 7
+    level['encoder_depths'] = [3, 3, 9, 3]
+    level['encoder_dims'] = [32, 64, 128, 256]
+    level['decoder_input'] = (16, 8, 80)
+    level['decoder_depths'] = [3, 3, 9, 3]
+    level['decoder_dims'] = [256, 128, 64, 16]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = ''
+    config['pyramid']['levels'].append(level)
+  elif res == 256:
+    level = dict()
+    level['resolution'] = 256
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 7
+    level['encoder_depths'] = [9, 9, 9, 9]
+    level['encoder_dims'] = [32, 64, 128, 256]
+    level['decoder_input'] = (8, 4, 128)
+    level['decoder_depths'] = [9, 9, 9, 9]
+    level['decoder_dims'] = [256, 128, 64, 16]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = ''#'checkpoints/up_3_d256_512_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  elif res == 128:
+    level = dict()
+    level['resolution'] = 128
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 5
+    level['encoder_depths'] = [16, 9, 9, 9]
+    level['encoder_dims'] = [64, 128, 128, 256]
+    level['decoder_input'] = (4, 2, 128)
+    level['decoder_depths'] = [16, 9, 9, 9]
+    level['decoder_dims'] = [256, 128, 128, 64]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = ''#'checkpoints/up_3_d128_256_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  elif res == 64:
+    level = dict()
+    level['resolution'] = 64
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 5
+    level['encoder_depths'] = [9, 9, 9]
+    level['encoder_dims'] = [64, 128, 256]
+    level['decoder_input'] = (4, 2, 256)
+    level['decoder_depths'] = [9, 9, 9]
+    level['decoder_dims'] = [256, 128, 64]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = 'checkpoints/64_512_cape_m'#checkpoints/up_3_d64_128_2'
+    level['upscale_enhancer_freeze'] = False
+    config['pyramid']['levels'].append(level)
+  elif res == 32:
+    level = dict()
+    level['resolution'] = 32
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 3
+    level['encoder_depths'] = [9, 9]
+    level['encoder_dims'] = [256, 256]
+    level['decoder_input'] = (4, 2, 256)
+    level['decoder_depths'] = [9, 9]
+    level['decoder_dims'] = [256, 256]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = ''#'checkpoints/up_3_d32_64_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  else:
+    raise NotImplementedError
+
+
+
+config['cond_ae'] = dict()
+config['cond_ae']['is_label'] = False
+config['cond_ae']['upscale_output'] = False
+config['cond_ae']['latent_dim'] = config['latent_dim']
+config['cond_ae']['kernel_size'] = 7
+config['cond_ae']['bottleneck_size'] = 64
+config['cond_ae']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae']['encoder_dims'] = [32, 64, 128, 256, config['cond_ae']['bottleneck_size']]
+config['cond_ae']['decoder_input'] = (8, 4, config['cond_ae']['bottleneck_size'])
+config['cond_ae']['decoder_depths'] = [3, 3, 9, 3, 3]
+config['cond_ae']['decoder_dims'] = [config['cond_ae']['bottleneck_size'], 256, 128, 64, 16]
+
+config['cond_ae_body'] = dict()
+config['cond_ae_body']['latent_dim'] = config['latent_dim']
+config['cond_ae_body']['kernel_size'] = 7
+config['cond_ae_body']['upscale_output'] = True
+config['cond_ae_body']['bottleneck_size'] = 512
+config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3, 3]
+config['cond_ae_body']['encoder_dims'] = [32, 64, 128, 256, 256, config['cond_ae_body']['bottleneck_size']]
+config['cond_ae_body']['decoder_input'] = (4, 2, config['cond_ae_body']['bottleneck_size'])
+config['cond_ae_body']['decoder_depths'] = [3, 3, 9, 3]
+config['cond_ae_body']['decoder_dims'] = [config['cond_ae_body']['bottleneck_size'], 128, 64, 16]

configs/config_abl_2_pyr_64_cape_f_03223_SS.py

@@ -0,0 +1,190 @@
+config = dict()
+
+config['gpu'] = '1'
+config['name'] = 'abl_2_pyr_64_cape_f_03223_SS'
+config['notes'] = ''
+config['tags'] = ['cape_f_2']
+
+# DON'T FORGET TO RESET SSIM LOSS USAGE
+# DON'T FORGET TO CHECK MODEL LOADING
+# DON'T FORGET TO CHECK IF STORE IN MEMORY
+# DON'T FORGET TO CHECK REG LOSS COEFF
+# DON'T FORGET TO LOAD FINAL UPSCALER UNLESS HAVE LEVEL 512
+# DON'T FORGET TO CHECK IF SKIRT
+# DON'T FORGET TO CHECK IF CAPE
+# DON'T FORGET TO CHECK IF USING DOWNSCALE NONE OR _2
+# DON'T FORGET TO CHECK UPSCALING and ENCODER CHECKPOINTS
+
+config['use_skinning'] =              True # and deshaping
+config['remove_global_orientation'] = True
+config['use_extra_losses'] =          True # edge, collision, 3d mesh
+config['use_offsets'] =               False
+config['use_norm_loss'] =             False
+
+config['pretrained_model'] = ''
+config['pretrained_model_iteration'] = ''
+
+config['pyramid'] = dict()
+config['pyramid']['resolutions'] = [64]
+config['pyramid']['pretrained_pyramid'] = ''
+config['pyramid']['final_upscale_enhancer'] = 'checkpoints/64_512_cape_f'
+
+config['train_dir'] = '/home/workspace/local-data/cape_processed/train/female'
+config['valid_dir'] = '/home/workspace/local-data/cape_processed/test/female'
+config['test_dir'] = '/home/workspace/local-data/cape_processed/test/female'
+config['batch_size'] = 4
+config['is_skirt'] = False
+config['is_cape'] = True
+
+config['frames_file'] = 'frames_03223_SS.txt'
+config['frames_file_valid'] = 'frames_03223_SS.txt'
+
+config['pyramid']['static_model'] = 'checkpoints/static_1_cape_f'
+config['pyramid']['body_model'] = 'checkpoints/body_1_cape_f'
+config['pyramid']['curvature_model'] = 'checkpoints/curv_1_cape_f'
+
+config['dataset_in_memory'] = True
+config['resolution'] = None
+config['latent_dim'] = 128
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['lr'] = 0.0001
+config['starting_balancing_coeff'] = 0.0001
+config['reg_loss_coeff'] = None # set to 0 if don't want to use edge+collision
+config['avg_garment_path'] = ''
+
+config['signature_dynamic'] = { 'id': (), 'frame_nr': (),
+                          'pose_shape': (72 + 11,),
+                          'uv_body_posed': (512, 256, 3),
+                          'mesh_body_posed': (config['smpl_res'], 3,),
+                          'mesh_posed': (config['smpl_res'], 3,),
+                          'mesh_unposed': (config['smpl_res'], 3,),
+                          'uv_unposed': (512, 256, 3),
+                          'uv_curvature': (512, 256, 7),
+                        }
+config['signature_static'] = { 'uv_static': (512, 256, 3),
+                              'mesh_mask': (config['smpl_res'],),
+                              'uv_mask': (512, 256,),
+                              }
+
+config['pyramid']['is_label'] = False
+config['pyramid']['latent_dim'] = config['latent_dim']
+config['pyramid']['add_noise'] = True
+config['pyramid']['noise_amount'] = 0.001
+
+config['pyramid']['loss_weights'] = [1,1,1,1,1]
+
+for res in config['pyramid']['resolutions']:
+  if res != 512:
+      res_modifier = '_' + str(res)
+      config['signature_dynamic']['uv_body_posed' + res_modifier] = (res, res//2, 3)
+      config['signature_dynamic']['uv_unposed' + res_modifier] = (res, res//2, 3)
+      config['signature_dynamic']['uv_curvature' + res_modifier] = (res, res//2, 7)
+      config['signature_static']['uv_static' + res_modifier] = (res, res//2, 3)
+      config['signature_static']['uv_mask' + res_modifier] = (res, res//2)
+
+
+
+config['pyramid']['freeze_encoders'] = True
+
+config['pyramid']['levels'] = []
+for res in config['pyramid']['resolutions']:
+  if res == 512:
+    level = dict()
+    level['resolution'] = 512
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 7
+    level['encoder_depths'] = [3, 3, 9, 3]
+    level['encoder_dims'] = [32, 64, 128, 256]
+    level['decoder_input'] = (16, 8, 80)
+    level['decoder_depths'] = [3, 3, 9, 3]
+    level['decoder_dims'] = [256, 128, 64, 16]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = ''
+    config['pyramid']['levels'].append(level)
+  elif res == 256:
+    level = dict()
+    level['resolution'] = 256
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 7
+    level['encoder_depths'] = [9, 9, 9, 9]
+    level['encoder_dims'] = [32, 64, 128, 256]
+    level['decoder_input'] = (8, 4, 128)
+    level['decoder_depths'] = [9, 9, 9, 9]
+    level['decoder_dims'] = [256, 128, 64, 16]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = ''#'checkpoints/up_3_d256_512_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  elif res == 128:
+    level = dict()
+    level['resolution'] = 128
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 5
+    level['encoder_depths'] = [16, 9, 9, 9]
+    level['encoder_dims'] = [64, 128, 128, 256]
+    level['decoder_input'] = (4, 2, 128)
+    level['decoder_depths'] = [16, 9, 9, 9]
+    level['decoder_dims'] = [256, 128, 128, 64]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = ''#'checkpoints/up_3_d128_256_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  elif res == 64:
+    level = dict()
+    level['resolution'] = 64
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 5
+    level['encoder_depths'] = [9, 9, 9]
+    level['encoder_dims'] = [64, 128, 256]
+    level['decoder_input'] = (4, 2, 256)
+    level['decoder_depths'] = [9, 9, 9]
+    level['decoder_dims'] = [256, 128, 64]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = ''#checkpoints/up_3_d64_128_2'
+    level['upscale_enhancer_freeze'] = False
+    config['pyramid']['levels'].append(level)
+  elif res == 32:
+    level = dict()
+    level['resolution'] = 32
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 3
+    level['encoder_depths'] = [9, 9]
+    level['encoder_dims'] = [256, 256]
+    level['decoder_input'] = (4, 2, 256)
+    level['decoder_depths'] = [9, 9]
+    level['decoder_dims'] = [256, 256]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = ''#'checkpoints/up_3_d32_64_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  else:
+    raise NotImplementedError
+
+
+
+config['cond_ae'] = dict()
+config['cond_ae']['is_label'] = False
+config['cond_ae']['upscale_output'] = False
+config['cond_ae']['latent_dim'] = config['latent_dim']
+config['cond_ae']['kernel_size'] = 7
+config['cond_ae']['bottleneck_size'] = 64
+config['cond_ae']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae']['encoder_dims'] = [32, 64, 128, 256, config['cond_ae']['bottleneck_size']]
+config['cond_ae']['decoder_input'] = (8, 4, config['cond_ae']['bottleneck_size'])
+config['cond_ae']['decoder_depths'] = [3, 3, 9, 3, 3]
+config['cond_ae']['decoder_dims'] = [config['cond_ae']['bottleneck_size'], 256, 128, 64, 16]
+
+config['cond_ae_body'] = dict()
+config['cond_ae_body']['latent_dim'] = config['latent_dim']
+config['cond_ae_body']['kernel_size'] = 7
+config['cond_ae_body']['upscale_output'] = True
+config['cond_ae_body']['bottleneck_size'] = 512
+config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3, 3]
+config['cond_ae_body']['encoder_dims'] = [32, 64, 128, 256, 256, config['cond_ae_body']['bottleneck_size']]
+config['cond_ae_body']['decoder_input'] = (4, 2, config['cond_ae_body']['bottleneck_size'])
+config['cond_ae_body']['decoder_depths'] = [3, 3, 9, 3]
+config['cond_ae_body']['decoder_dims'] = [config['cond_ae_body']['bottleneck_size'], 128, 64, 16]

configs/config_abl_2_pyr_64_cape_m_00032_SL.py

@@ -0,0 +1,190 @@
+config = dict()
+
+config['gpu'] = '2'
+config['name'] = 'abl_2_pyr_64_cape_m_00032_SL'
+config['notes'] = ''
+config['tags'] = ['cape_m_2']
+
+# DON'T FORGET TO RESET SSIM LOSS USAGE
+# DON'T FORGET TO CHECK MODEL LOADING
+# DON'T FORGET TO CHECK IF STORE IN MEMORY
+# DON'T FORGET TO CHECK REG LOSS COEFF
+# DON'T FORGET TO LOAD FINAL UPSCALER UNLESS HAVE LEVEL 512
+# DON'T FORGET TO CHECK IF SKIRT
+# DON'T FORGET TO CHECK IF CAPE
+# DON'T FORGET TO CHECK IF USING DOWNSCALE NONE OR _2
+# DON'T FORGET TO CHECK UPSCALING and ENCODER CHECKPOINTS
+
+config['use_skinning'] =              True # and deshaping
+config['remove_global_orientation'] = True
+config['use_extra_losses'] =          True # edge, collision, 3d mesh
+config['use_offsets'] =               False
+config['use_norm_loss'] =             False
+
+config['pretrained_model'] = ''
+config['pretrained_model_iteration'] = ''
+
+config['pyramid'] = dict()
+config['pyramid']['resolutions'] = [64]
+config['pyramid']['pretrained_pyramid'] = ''
+config['pyramid']['final_upscale_enhancer'] = 'checkpoints/64_512_cape_m'
+
+config['train_dir'] = '/home/workspace/local-data/cape_processed/train/male'
+config['valid_dir'] = '/home/workspace/local-data/cape_processed/train/male'
+config['test_dir'] = '/home/workspace/local-data/cape_processed/train/male'
+config['batch_size'] = 4
+config['is_skirt'] = False
+config['is_cape'] = True
+
+config['frames_file'] = 'frames_00032_SL_train.txt'
+config['frames_file_valid'] = 'frames_00032_SL_test.txt'
+
+config['pyramid']['static_model'] = 'checkpoints/static_1_cape_m'
+config['pyramid']['body_model'] = 'checkpoints/body_1_cape_m'
+config['pyramid']['curvature_model'] = 'checkpoints/curv_1_cape_m'
+
+config['dataset_in_memory'] = True
+config['resolution'] = None
+config['latent_dim'] = 128
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['lr'] = 0.0001
+config['starting_balancing_coeff'] = 0.0001
+config['reg_loss_coeff'] = None # set to 0 if don't want to use edge+collision
+config['avg_garment_path'] = ''
+
+config['signature_dynamic'] = { 'id': (), 'frame_nr': (),
+                          'pose_shape': (72 + 11,),
+                          'uv_body_posed': (512, 256, 3),
+                          'mesh_body_posed': (config['smpl_res'], 3,),
+                          'mesh_posed': (config['smpl_res'], 3,),
+                          'mesh_unposed': (config['smpl_res'], 3,),
+                          'uv_unposed': (512, 256, 3),
+                          'uv_curvature': (512, 256, 7),
+                        }
+config['signature_static'] = { 'uv_static': (512, 256, 3),
+                              'mesh_mask': (config['smpl_res'],),
+                              'uv_mask': (512, 256,),
+                              }
+
+config['pyramid']['is_label'] = False
+config['pyramid']['latent_dim'] = config['latent_dim']
+config['pyramid']['add_noise'] = True
+config['pyramid']['noise_amount'] = 0.001
+
+config['pyramid']['loss_weights'] = [1,1,1,1,1]
+
+for res in config['pyramid']['resolutions']:
+  if res != 512:
+      res_modifier = '_' + str(res)
+      config['signature_dynamic']['uv_body_posed' + res_modifier] = (res, res//2, 3)
+      config['signature_dynamic']['uv_unposed' + res_modifier] = (res, res//2, 3)
+      config['signature_dynamic']['uv_curvature' + res_modifier] = (res, res//2, 7)
+      config['signature_static']['uv_static' + res_modifier] = (res, res//2, 3)
+      config['signature_static']['uv_mask' + res_modifier] = (res, res//2)
+
+
+
+config['pyramid']['freeze_encoders'] = True
+
+config['pyramid']['levels'] = []
+for res in config['pyramid']['resolutions']:
+  if res == 512:
+    level = dict()
+    level['resolution'] = 512
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 7
+    level['encoder_depths'] = [3, 3, 9, 3]
+    level['encoder_dims'] = [32, 64, 128, 256]
+    level['decoder_input'] = (16, 8, 80)
+    level['decoder_depths'] = [3, 3, 9, 3]
+    level['decoder_dims'] = [256, 128, 64, 16]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = ''
+    config['pyramid']['levels'].append(level)
+  elif res == 256:
+    level = dict()
+    level['resolution'] = 256
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 7
+    level['encoder_depths'] = [9, 9, 9, 9]
+    level['encoder_dims'] = [32, 64, 128, 256]
+    level['decoder_input'] = (8, 4, 128)
+    level['decoder_depths'] = [9, 9, 9, 9]
+    level['decoder_dims'] = [256, 128, 64, 16]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = ''#'checkpoints/up_3_d256_512_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  elif res == 128:
+    level = dict()
+    level['resolution'] = 128
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 5
+    level['encoder_depths'] = [16, 9, 9, 9]
+    level['encoder_dims'] = [64, 128, 128, 256]
+    level['decoder_input'] = (4, 2, 128)
+    level['decoder_depths'] = [16, 9, 9, 9]
+    level['decoder_dims'] = [256, 128, 128, 64]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = ''#'checkpoints/up_3_d128_256_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  elif res == 64:
+    level = dict()
+    level['resolution'] = 64
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 5
+    level['encoder_depths'] = [9, 9, 9]
+    level['encoder_dims'] = [64, 128, 256]
+    level['decoder_input'] = (4, 2, 256)
+    level['decoder_depths'] = [9, 9, 9]
+    level['decoder_dims'] = [256, 128, 64]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = ''#checkpoints/up_3_d64_128_2'
+    level['upscale_enhancer_freeze'] = False
+    config['pyramid']['levels'].append(level)
+  elif res == 32:
+    level = dict()
+    level['resolution'] = 32
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 3
+    level['encoder_depths'] = [9, 9]
+    level['encoder_dims'] = [256, 256]
+    level['decoder_input'] = (4, 2, 256)
+    level['decoder_depths'] = [9, 9]
+    level['decoder_dims'] = [256, 256]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = ''#'checkpoints/up_3_d32_64_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  else:
+    raise NotImplementedError
+
+
+
+config['cond_ae'] = dict()
+config['cond_ae']['is_label'] = False
+config['cond_ae']['upscale_output'] = False
+config['cond_ae']['latent_dim'] = config['latent_dim']
+config['cond_ae']['kernel_size'] = 7
+config['cond_ae']['bottleneck_size'] = 64
+config['cond_ae']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae']['encoder_dims'] = [32, 64, 128, 256, config['cond_ae']['bottleneck_size']]
+config['cond_ae']['decoder_input'] = (8, 4, config['cond_ae']['bottleneck_size'])
+config['cond_ae']['decoder_depths'] = [3, 3, 9, 3, 3]
+config['cond_ae']['decoder_dims'] = [config['cond_ae']['bottleneck_size'], 256, 128, 64, 16]
+
+config['cond_ae_body'] = dict()
+config['cond_ae_body']['latent_dim'] = config['latent_dim']
+config['cond_ae_body']['kernel_size'] = 7
+config['cond_ae_body']['upscale_output'] = True
+config['cond_ae_body']['bottleneck_size'] = 512
+config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3, 3]
+config['cond_ae_body']['encoder_dims'] = [32, 64, 128, 256, 256, config['cond_ae_body']['bottleneck_size']]
+config['cond_ae_body']['decoder_input'] = (4, 2, config['cond_ae_body']['bottleneck_size'])
+config['cond_ae_body']['decoder_depths'] = [3, 3, 9, 3]
+config['cond_ae_body']['decoder_dims'] = [config['cond_ae_body']['bottleneck_size'], 128, 64, 16]

configs/config_abl_50p_skin.py

@@ -0,0 +1,118 @@
+config = dict()
+
+# DON'T FORGET TO RESET SSIM LOSS USAGE
+# DON'T FORGET TO CHECK MODEL LOADING
+# DON'T FORGET TO CHECK IF STORE IN MEMORY
+# DON'T FORGET TO CHECK REG LOSS COEFF
+# DON'T FORGET TO LOAD FINAL UPSCALER UNLESS HAVE LEVEL 512
+# DON'T FORGET TO CHECK AVG GARMENT USED
+
+config['use_skinning'] =              False # and deshaping
+config['remove_global_orientation'] = True
+config['use_curvature'] =             True
+config['use_extra_losses'] =          True # edge, collision, 3d mesh
+config['use_offsets'] =               True
+
+config['use_norm_loss'] =             False
+
+ablation_count = [config['use_skinning'], config['remove_global_orientation'], config['use_curvature'], config['use_extra_losses'], config['use_offsets']].count(False)
+assert ablation_count == 1 or ablation_count == 0
+
+config['gpu'] = '6'
+config['name'] = 'abl_50p_skin'
+config['notes'] = ''
+
+config['avg_garment_path'] = '' #custom_data/train_no_shape_avg_uv.npy # train_no_shape_avg_uv  train_no_shape_one_hot_avg_uv
+if config['use_skinning']:
+  assert (config['use_offsets'] and config['avg_garment_path'] != '') or (not config['use_offsets'] and config['avg_garment_path'] == '')
+
+config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
+
+config['batch_size'] = 4
+config['resolution'] = 512
+config['is_skirt'] = False
+config['dataset_in_memory'] = False
+config['is_cape'] = False
+
+config['train_dir'] = '/home/workspace/local-data/train'
+config['valid_dir'] = '/home/workspace/local-data/valid'
+config['test_dir'] = '/home/workspace/data/dataset/test_ready'
+config['frames_file'] = 'frames_50p.txt'
+
+config['static_model'] = 'checkpoints/static_1'
+config['body_model'] = 'checkpoints/body_1'
+config['curvature_model'] = 'checkpoints/curv_1'
+
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+config['pretrained_model'] = ''
+
+###############    
+config['lr'] = 0.0001#                 0.0001
+###############
+config['starting_balancing_coeff'] = 0.0001
+
+config['signature_dynamic'] = { 'id': (), 'frame_nr': (),
+                          'pose_shape': (72 + 11,),
+                          'uv_body_posed': (512, 256, 3),
+                          'mesh_body_posed': (14475, 3,),
+                          'mesh_posed': (14475, 3,),
+                          'uv_curvature': (512, 256, 7),
+                        }
+config['signature_static'] = { 'uv_static': (512, 256, 3),
+                              'mesh_mask': (14475,),
+                              'uv_mask': (512, 256,),
+                              }
+
+if config['use_skinning']:
+  config['signature_dynamic']['uv_unposed'] = (512, 256, 3)
+  config['signature_dynamic']['mesh_unposed'] = (14475, 3,)
+else:
+  config['signature_dynamic']['uv_posed'] = (512, 256, 3)
+
+
+config['upscale_enhancer'] = ''
+config['freeze_encoders'] = True
+config['use_one_hot'] = False
+
+config['cvae'] = dict()
+config['cvae']['is_label'] = False
+config['cvae']['latent_dim'] = config['latent_dim']
+config['cvae']['add_noise'] = True
+config['cvae']['noise_amount'] = 0.025
+
+config['cvae']['kernel_size'] = 7
+config['cvae']['encoder_depths'] = [3, 3, 9, 3]
+config['cvae']['encoder_dims'] = [32, 64, 128, 256]
+config['cvae']['decoder_input'] = (16, 8, 80)
+config['cvae']['decoder_depths'] = [3, 3, 9, 3]
+config['cvae']['decoder_dims'] = [256, 128, 64, 16]
+
+
+config['cond_ae'] = dict()
+config['cond_ae']['is_label'] = False
+config['cond_ae']['upscale_output'] = False
+config['cond_ae']['latent_dim'] = config['latent_dim']
+config['cond_ae']['kernel_size'] = 7
+config['cond_ae']['bottleneck_size'] = 64
+config['cond_ae']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae']['encoder_dims'] = [32, 64, 128, 256, config['cond_ae']['bottleneck_size']]
+config['cond_ae']['decoder_input'] = (8, 4, config['cond_ae']['bottleneck_size'])
+config['cond_ae']['decoder_depths'] = [3, 3, 9, 3, 3]
+config['cond_ae']['decoder_dims'] = [config['cond_ae']['bottleneck_size'], 256, 128, 64, 16]
+
+config['cond_ae_body'] = dict()
+config['cond_ae_body']['latent_dim'] = config['latent_dim']
+config['cond_ae_body']['kernel_size'] = 7
+config['cond_ae_body']['upscale_output'] = True
+config['cond_ae_body']['bottleneck_size'] = 512
+config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3, 3]
+config['cond_ae_body']['encoder_dims'] = [32, 64, 128, 256, 256, config['cond_ae_body']['bottleneck_size']]
+config['cond_ae_body']['decoder_input'] = (4, 2, config['cond_ae_body']['bottleneck_size'])
+config['cond_ae_body']['decoder_depths'] = [3, 3, 9, 3]
+config['cond_ae_body']['decoder_dims'] = [config['cond_ae_body']['bottleneck_size'], 128, 64, 16]

configs/config_abl_all_post.py

@@ -0,0 +1,126 @@
+config = dict()
+
+# DON'T FORGET TO RESET SSIM LOSS USAGE
+# DON'T FORGET TO CHECK MODEL LOADING
+# DON'T FORGET TO CHECK IF STORE IN MEMORY
+# DON'T FORGET TO CHECK REG LOSS COEFF
+# DON'T FORGET TO LOAD FINAL UPSCALER UNLESS HAVE LEVEL 512
+# DON'T FORGET TO CHECK AVG GARMENT USED
+
+config['use_skinning'] =              True # and deshaping
+config['remove_global_orientation'] = True
+config['use_curvature'] =             True
+config['use_extra_losses'] =          True # edge, collision, 3d mesh
+config['use_offsets'] =               True
+
+config['use_norm_loss'] =             True
+config['only_curvatures'] =           False
+
+ablation_count = [config['use_skinning'], config['remove_global_orientation'], config['use_curvature'], config['use_extra_losses'], config['use_offsets']].count(False)
+assert ablation_count == 1 or ablation_count == 0
+
+config['gpu'] = '1'
+config['name'] = 'al_all_post'
+config['notes'] = ''
+
+config['avg_garment_path'] = 'custom_data/train_no_shape_avg_uv.npy' #custom_data/train_no_shape_avg_uv.npy # train_no_shape_avg_uv  train_no_shape_one_hot_avg_uv
+if config['use_skinning']:
+  assert (config['use_offsets'] and config['avg_garment_path'] != '') or (not config['use_offsets'] and config['avg_garment_path'] == '')
+
+config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
+
+config['batch_size'] = 4
+config['resolution'] = 512
+config['is_skirt'] = False
+config['dataset_in_memory'] = False
+config['is_cape'] = False
+
+config['train_dir'] = '/home/workspace/local-data/train'
+config['valid_dir'] = '/home/workspace/local-data/valid'
+config['test_dir'] = '/home/workspace/data/dataset/test_ready'
+
+# config['frames_file'] = 'frames_dress_30k.txt'
+# config['frames_file_valid'] = 'frames_dress_30k.txt'
+# config['frames_file_test'] = 'frames_dress_30k.txt'
+
+config['static_model'] = 'checkpoints/post/static_post_2'
+config['body_model'] = 'checkpoints/post/body_post_2'
+config['curvature_model'] = 'checkpoints/post/curv_post_1'
+
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+
+config['pretrained_model'] = ''
+
+###############    
+config['lr'] = 0.0001#                 0.0001
+###############
+config['starting_balancing_coeff'] = 0.001
+
+config['signature_dynamic'] = { 'id': (), 'frame_nr': (),
+                          'pose_shape': (72 + 11,),
+                          'uv_body_posed': (512, 256, 3),
+                          'mesh_body_posed': (config['smpl_res'], 3,),
+                          'mesh_posed': (config['smpl_res'], 3,),
+                          'uv_curvature': (512, 256, 7),
+                        }
+config['signature_static'] = { 'uv_static': (512, 256, 3),
+                              'mesh_mask': (config['smpl_res'],),
+                              'uv_mask': (512, 256,),
+                              }
+
+if config['use_skinning']:
+  config['signature_dynamic']['uv_unposed'] = (512, 256, 3)
+  config['signature_dynamic']['mesh_unposed'] = (config['smpl_res'], 3,)
+else:
+  config['signature_dynamic']['uv_posed'] = (512, 256, 3)
+
+
+config['upscale_enhancer'] = ''
+config['freeze_encoders'] = True
+config['use_one_hot'] = False
+
+config['cvae'] = dict()
+config['cvae']['is_label'] = False
+config['cvae']['add_noise'] = True
+config['cvae']['noise_amount'] = 0.001
+
+config['cvae']['kernel_size'] = 7
+config['cvae']['encoder_depths'] = [3, 3, 3, 9, 3]
+config['cvae']['encoder_dims'] = [32, 64, 128, 256, 64]
+config['cvae']['decoder_depths'] = config['cvae']['encoder_depths']
+config['cvae']['decoder_dims'] = config['cvae']['encoder_dims']
+
+
+config['cond_ae_static'] = dict()
+config['cond_ae_static']['is_label'] = False
+config['cond_ae_static']['kernel_size'] = 7
+config['cond_ae_static']['bottleneck_size'] = 8
+config['cond_ae_static']['encoder_depths'] = [3, 3, 3, 9, 3]
+config['cond_ae_static']['encoder_dims'] = [16, 32, 64, 64, config['cond_ae_static']['bottleneck_size']]
+config['cond_ae_static']['decoder_depths'] = config['cond_ae_static']['encoder_depths']
+config['cond_ae_static']['decoder_dims'] = config['cond_ae_static']['encoder_dims']
+config['cond_ae_static']['activation'] = 'tanh'
+
+config['cond_ae_curv'] = dict()
+config['cond_ae_curv']['is_label'] = True
+config['cond_ae_curv']['kernel_size'] = 7
+config['cond_ae_curv']['bottleneck_size'] = 40
+config['cond_ae_curv']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae_curv']['encoder_dims'] = [16, 32, 64, 128, config['cond_ae_curv']['bottleneck_size']]
+config['cond_ae_curv']['decoder_depths'] = config['cond_ae_curv']['encoder_depths']
+config['cond_ae_curv']['decoder_dims'] = config['cond_ae_curv']['encoder_dims']
+config['cond_ae_curv']['activation'] = 'tanh'
+
+config['cond_ae_body'] = dict()
+config['cond_ae_body']['is_label'] = False
+config['cond_ae_body']['kernel_size'] = 7
+config['cond_ae_body']['bottleneck_size'] = 16
+config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae_body']['encoder_dims'] = [16, 32, 64, 128, config['cond_ae_body']['bottleneck_size']]
+config['cond_ae_body']['decoder_depths'] = config['cond_ae_body']['encoder_depths']
+config['cond_ae_body']['decoder_dims'] = config['cond_ae_body']['encoder_dims']
+config['cond_ae_body']['activation'] = 'tanh'

configs/config_abl_norm.py

@@ -18,7 +18,7 @@ config['use_norm_loss'] =             True
 ablation_count = [config['use_skinning'], config['remove_global_orientation'], config['use_curvature'], config['use_extra_losses'], config['use_offsets']].count(False)
 assert ablation_count == 1 or ablation_count == 0
 
-config['gpu'] = '2'
+config['gpu'] = '6'
 config['name'] = 'abl_norm'
 config['notes'] = ''
 

configs/config_abl_norm_skirt.py

@@ -18,7 +18,7 @@ config['use_norm_loss'] =             True
 ablation_count = [config['use_skinning'], config['remove_global_orientation'], config['use_curvature'], config['use_extra_losses'], config['use_offsets']].count(False)
 assert ablation_count == 1 or ablation_count == 0
 
-config['gpu'] = '0'
+config['gpu'] = '4'
 config['name'] = 'abl_norm_skirt'
 config['notes'] = ''
 

configs/config_abl_overfit.py

@@ -0,0 +1,121 @@
+config = dict()
+
+# DON'T FORGET TO RESET SSIM LOSS USAGE
+# DON'T FORGET TO CHECK MODEL LOADING
+# DON'T FORGET TO CHECK IF STORE IN MEMORY
+# DON'T FORGET TO CHECK REG LOSS COEFF
+# DON'T FORGET TO LOAD FINAL UPSCALER UNLESS HAVE LEVEL 512
+# DON'T FORGET TO CHECK AVG GARMENT USED
+
+config['use_skinning'] =              True # and deshaping
+config['remove_global_orientation'] = True
+config['use_curvature'] =             True
+config['use_extra_losses'] =          True # edge, collision, 3d mesh
+config['use_offsets'] =               True
+
+config['use_norm_loss'] =             True
+
+ablation_count = [config['use_skinning'], config['remove_global_orientation'], config['use_curvature'], config['use_extra_losses'], config['use_offsets']].count(False)
+assert ablation_count == 1 or ablation_count == 0
+
+config['gpu'] = '0'
+config['name'] = 'abl_overfit'
+config['notes'] = ''
+
+config['avg_garment_path'] = 'custom_data/train_00761_no_shape_avg_uv.npy' #custom_data/train_no_shape_avg_uv.npy # train_no_shape_avg_uv  train_no_shape_one_hot_avg_uv
+if config['use_skinning']:
+  assert (config['use_offsets'] and config['avg_garment_path'] != '') or (not config['use_offsets'] and config['avg_garment_path'] == '')
+
+config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
+
+config['batch_size'] = 4
+config['resolution'] = 512
+config['is_skirt'] = False
+config['dataset_in_memory'] = True
+config['is_cape'] = False
+
+config['train_dir'] = '/home/workspace/local-data/train'
+config['valid_dir'] = '/home/workspace/local-data/train'
+config['test_dir'] = '/home/workspace/local-data/train'
+
+config['frames_file'] = 'frames_00761.txt'
+config['frames_file_valid'] = 'frames_00761.txt'
+config['frames_file_test'] = 'frames_00761.txt'
+
+config['static_model'] = 'checkpoints/static_1'
+config['body_model'] = 'checkpoints/body_1'
+config['curvature_model'] = 'checkpoints/curv_1'
+
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+config['pretrained_model'] = ''
+
+###############    
+config['lr'] = 0.0001#                 0.0001
+###############
+config['starting_balancing_coeff'] = 0.0001
+
+config['signature_dynamic'] = { 'id': (), 'frame_nr': (),
+                          'pose_shape': (72 + 11,),
+                          'uv_body_posed': (512, 256, 3),
+                          'mesh_body_posed': (14475, 3,),
+                          'mesh_posed': (14475, 3,),
+                          'uv_curvature': (512, 256, 7),
+                        }
+config['signature_static'] = { 'uv_static': (512, 256, 3),
+                              'mesh_mask': (14475,),
+                              'uv_mask': (512, 256,),
+                              }
+
+if config['use_skinning']:
+  config['signature_dynamic']['uv_unposed'] = (512, 256, 3)
+  config['signature_dynamic']['mesh_unposed'] = (14475, 3,)
+else:
+  config['signature_dynamic']['uv_posed'] = (512, 256, 3)
+
+
+config['upscale_enhancer'] = ''
+config['freeze_encoders'] = True
+config['use_one_hot'] = False
+
+config['cvae'] = dict()
+config['cvae']['is_label'] = False
+config['cvae']['latent_dim'] = config['latent_dim']
+config['cvae']['add_noise'] = True
+config['cvae']['noise_amount'] = 0.025
+
+config['cvae']['kernel_size'] = 7
+config['cvae']['encoder_depths'] = [3, 3, 9, 3]
+config['cvae']['encoder_dims'] = [32, 64, 128, 256]
+config['cvae']['decoder_input'] = (16, 8, 80)
+config['cvae']['decoder_depths'] = [3, 3, 9, 3]
+config['cvae']['decoder_dims'] = [256, 128, 64, 16]
+
+
+config['cond_ae'] = dict()
+config['cond_ae']['is_label'] = False
+config['cond_ae']['upscale_output'] = False
+config['cond_ae']['latent_dim'] = config['latent_dim']
+config['cond_ae']['kernel_size'] = 7
+config['cond_ae']['bottleneck_size'] = 64
+config['cond_ae']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae']['encoder_dims'] = [32, 64, 128, 256, config['cond_ae']['bottleneck_size']]
+config['cond_ae']['decoder_input'] = (8, 4, config['cond_ae']['bottleneck_size'])
+config['cond_ae']['decoder_depths'] = [3, 3, 9, 3, 3]
+config['cond_ae']['decoder_dims'] = [config['cond_ae']['bottleneck_size'], 256, 128, 64, 16]
+
+config['cond_ae_body'] = dict()
+config['cond_ae_body']['latent_dim'] = config['latent_dim']
+config['cond_ae_body']['kernel_size'] = 7
+config['cond_ae_body']['upscale_output'] = True
+config['cond_ae_body']['bottleneck_size'] = 512
+config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3, 3]
+config['cond_ae_body']['encoder_dims'] = [32, 64, 128, 256, 256, config['cond_ae_body']['bottleneck_size']]
+config['cond_ae_body']['decoder_input'] = (4, 2, config['cond_ae_body']['bottleneck_size'])
+config['cond_ae_body']['decoder_depths'] = [3, 3, 9, 3]
+config['cond_ae_body']['decoder_dims'] = [config['cond_ae_body']['bottleneck_size'], 128, 64, 16]

configs/config_abl_overfit_dress.py

@@ -0,0 +1,122 @@
+config = dict()
+
+# DON'T FORGET TO RESET SSIM LOSS USAGE
+# DON'T FORGET TO CHECK MODEL LOADING
+# DON'T FORGET TO CHECK IF STORE IN MEMORY
+# DON'T FORGET TO CHECK REG LOSS COEFF
+# DON'T FORGET TO LOAD FINAL UPSCALER UNLESS HAVE LEVEL 512
+# DON'T FORGET TO CHECK AVG GARMENT USED
+
+config['use_skinning'] =              True # and deshaping
+config['remove_global_orientation'] = True
+config['use_curvature'] =             True
+config['use_extra_losses'] =          True # edge, collision, 3d mesh
+config['use_offsets'] =               True
+
+config['use_norm_loss'] =             True
+config['only_curvatures'] =           True
+
+ablation_count = [config['use_skinning'], config['remove_global_orientation'], config['use_curvature'], config['use_extra_losses'], config['use_offsets']].count(False)
+assert ablation_count == 1 or ablation_count == 0
+
+config['gpu'] = '4'
+config['name'] = 'abl_overfit_dress'
+config['notes'] = 'remove other conditions from input as well'
+
+config['avg_garment_path'] = 'custom_data/train_dress_30k_no_shape_avg_uv.npy' #custom_data/train_no_shape_avg_uv.npy # train_no_shape_avg_uv  train_no_shape_one_hot_avg_uv
+if config['use_skinning']:
+  assert (config['use_offsets'] and config['avg_garment_path'] != '') or (not config['use_offsets'] and config['avg_garment_path'] == '')
+
+config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
+
+config['batch_size'] = 4
+config['resolution'] = 512
+config['is_skirt'] = True
+config['dataset_in_memory'] = True
+config['is_cape'] = False
+
+config['train_dir'] = '/home/workspace/local-data/train_skirt'
+config['valid_dir'] = '/home/workspace/local-data/train_skirt'
+config['test_dir'] = '/home/workspace/local-data/train_skirt'
+
+config['frames_file'] = 'frames_dress_30k.txt'
+config['frames_file_valid'] = 'frames_dress_30k.txt'
+config['frames_file_test'] = 'frames_dress_30k.txt'
+
+config['static_model'] = 'checkpoints/static_1_skirt'
+config['body_model'] = 'checkpoints/body_1_skirt_2_2'
+config['curvature_model'] = 'checkpoints/curv_1_skirt'
+
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+config['pretrained_model'] = ''
+
+###############    
+config['lr'] = 0.0001#                 0.0001
+###############
+config['starting_balancing_coeff'] = 0.01
+
+config['signature_dynamic'] = { 'id': (), 'frame_nr': (),
+                          'pose_shape': (72 + 11,),
+                          'uv_body_posed': (512, 256, 3),
+                          'mesh_body_posed': (14475, 3,),
+                          'mesh_posed': (14475, 3,),
+                          'uv_curvature': (512, 256, 7),
+                        }
+config['signature_static'] = { 'uv_static': (512, 256, 3),
+                              'mesh_mask': (14475,),
+                              'uv_mask': (512, 256,),
+                              }
+
+if config['use_skinning']:
+  config['signature_dynamic']['uv_unposed'] = (512, 256, 3)
+  config['signature_dynamic']['mesh_unposed'] = (14475, 3,)
+else:
+  config['signature_dynamic']['uv_posed'] = (512, 256, 3)
+
+
+config['upscale_enhancer'] = ''
+config['freeze_encoders'] = True
+config['use_one_hot'] = False
+
+config['cvae'] = dict()
+config['cvae']['is_label'] = False
+config['cvae']['latent_dim'] = config['latent_dim']
+config['cvae']['add_noise'] = True
+config['cvae']['noise_amount'] = 0.025
+
+config['cvae']['kernel_size'] = 7
+config['cvae']['encoder_depths'] = [3, 3, 9, 3]
+config['cvae']['encoder_dims'] = [32, 64, 128, 256]
+config['cvae']['decoder_input'] = (16, 8, 80)
+config['cvae']['decoder_depths'] = [3, 3, 9, 3]
+config['cvae']['decoder_dims'] = [256, 128, 64, 16]
+
+
+config['cond_ae'] = dict()
+config['cond_ae']['is_label'] = False
+config['cond_ae']['upscale_output'] = False
+config['cond_ae']['latent_dim'] = config['latent_dim']
+config['cond_ae']['kernel_size'] = 7
+config['cond_ae']['bottleneck_size'] = 64
+config['cond_ae']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae']['encoder_dims'] = [32, 64, 128, 256, config['cond_ae']['bottleneck_size']]
+config['cond_ae']['decoder_input'] = (8, 4, config['cond_ae']['bottleneck_size'])
+config['cond_ae']['decoder_depths'] = [3, 3, 9, 3, 3]
+config['cond_ae']['decoder_dims'] = [config['cond_ae']['bottleneck_size'], 256, 128, 64, 16]
+
+config['cond_ae_body'] = dict()
+config['cond_ae_body']['latent_dim'] = config['latent_dim']
+config['cond_ae_body']['kernel_size'] = 7
+config['cond_ae_body']['upscale_output'] = True
+config['cond_ae_body']['bottleneck_size'] = 512
+config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3, 3]
+config['cond_ae_body']['encoder_dims'] = [32, 64, 128, 256, 256, config['cond_ae_body']['bottleneck_size']]
+config['cond_ae_body']['decoder_input'] = (4, 2, config['cond_ae_body']['bottleneck_size'])
+config['cond_ae_body']['decoder_depths'] = [3, 3, 9, 3]
+config['cond_ae_body']['decoder_dims'] = [config['cond_ae_body']['bottleneck_size'], 128, 64, 16]

configs/config_abl_overfit_dress_var.py

@@ -0,0 +1,123 @@
+config = dict()
+
+# DON'T FORGET TO RESET SSIM LOSS USAGE
+# DON'T FORGET TO CHECK MODEL LOADING
+# DON'T FORGET TO CHECK IF STORE IN MEMORY
+# DON'T FORGET TO CHECK REG LOSS COEFF
+# DON'T FORGET TO LOAD FINAL UPSCALER UNLESS HAVE LEVEL 512
+# DON'T FORGET TO CHECK AVG GARMENT USED
+
+config['use_skinning'] =              True # and deshaping
+config['remove_global_orientation'] = True
+config['use_curvature'] =             True
+config['use_extra_losses'] =          True # edge, collision, 3d mesh
+config['use_offsets'] =               True
+
+config['use_norm_loss'] =             True
+config['only_curvatures'] =           True
+
+ablation_count = [config['use_skinning'], config['remove_global_orientation'], config['use_curvature'], config['use_extra_losses'], config['use_offsets']].count(False)
+assert ablation_count == 1 or ablation_count == 0
+
+config['gpu'] = '0'
+config['name'] = 'abl_overfit_dress_var'
+config['notes'] = 'remove other conditions from input as well'
+
+config['avg_garment_path'] = 'custom_data/train_dress_30k_no_shape_avg_uv.npy' #custom_data/train_no_shape_avg_uv.npy # train_no_shape_avg_uv  train_no_shape_one_hot_avg_uv
+if config['use_skinning']:
+  assert (config['use_offsets'] and config['avg_garment_path'] != '') or (not config['use_offsets'] and config['avg_garment_path'] == '')
+
+config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
+
+config['batch_size'] = 4
+config['resolution'] = 512
+config['is_skirt'] = True
+config['dataset_in_memory'] = True
+config['is_cape'] = False
+
+config['train_dir'] = '/home/workspace/local-data/train_skirt'
+config['valid_dir'] = '/home/workspace/local-data/train_skirt'
+config['test_dir'] = '/home/workspace/local-data/train_skirt'
+
+config['frames_file'] = 'frames_dress_30k.txt'
+config['frames_file_valid'] = 'frames_dress_30k.txt'
+config['frames_file_test'] = 'frames_dress_30k.txt'
+
+config['static_model'] = 'checkpoints/static_1_skirt'
+config['body_model'] = 'checkpoints/body_1_skirt_2_2'
+config['curvature_model'] = 'checkpoints/curv_1_skirt'
+
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+config['pretrained_model'] = 'checkpoints/abl_overfit_dress'
+
+###############    
+config['lr'] = 0.0001#                 0.0001
+###############
+config['starting_balancing_coeff'] = 1.0
+config['balancing_activation_limit'] = 0.013
+
+config['signature_dynamic'] = { 'id': (), 'frame_nr': (),
+                          'pose_shape': (72 + 11,),
+                          'uv_body_posed': (512, 256, 3),
+                          'mesh_body_posed': (14475, 3,),
+                          'mesh_posed': (14475, 3,),
+                          'uv_curvature': (512, 256, 7),
+                        }
+config['signature_static'] = { 'uv_static': (512, 256, 3),
+                              'mesh_mask': (14475,),
+                              'uv_mask': (512, 256,),
+                              }
+
+if config['use_skinning']:
+  config['signature_dynamic']['uv_unposed'] = (512, 256, 3)
+  config['signature_dynamic']['mesh_unposed'] = (14475, 3,)
+else:
+  config['signature_dynamic']['uv_posed'] = (512, 256, 3)
+
+
+config['upscale_enhancer'] = ''
+config['freeze_encoders'] = True
+config['use_one_hot'] = False
+
+config['cvae'] = dict()
+config['cvae']['is_label'] = False
+config['cvae']['latent_dim'] = config['latent_dim']
+config['cvae']['add_noise'] = True
+config['cvae']['noise_amount'] = 0.025
+
+config['cvae']['kernel_size'] = 7
+config['cvae']['encoder_depths'] = [3, 3, 9, 3]
+config['cvae']['encoder_dims'] = [32, 64, 128, 256]
+config['cvae']['decoder_input'] = (16, 8, 80)
+config['cvae']['decoder_depths'] = [3, 3, 9, 3]
+config['cvae']['decoder_dims'] = [256, 128, 64, 16]
+
+
+config['cond_ae'] = dict()
+config['cond_ae']['is_label'] = False
+config['cond_ae']['upscale_output'] = False
+config['cond_ae']['latent_dim'] = config['latent_dim']
+config['cond_ae']['kernel_size'] = 7
+config['cond_ae']['bottleneck_size'] = 64
+config['cond_ae']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae']['encoder_dims'] = [32, 64, 128, 256, config['cond_ae']['bottleneck_size']]
+config['cond_ae']['decoder_input'] = (8, 4, config['cond_ae']['bottleneck_size'])
+config['cond_ae']['decoder_depths'] = [3, 3, 9, 3, 3]
+config['cond_ae']['decoder_dims'] = [config['cond_ae']['bottleneck_size'], 256, 128, 64, 16]
+
+config['cond_ae_body'] = dict()
+config['cond_ae_body']['latent_dim'] = config['latent_dim']
+config['cond_ae_body']['kernel_size'] = 7
+config['cond_ae_body']['upscale_output'] = True
+config['cond_ae_body']['bottleneck_size'] = 512
+config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3, 3]
+config['cond_ae_body']['encoder_dims'] = [32, 64, 128, 256, 256, config['cond_ae_body']['bottleneck_size']]
+config['cond_ae_body']['decoder_input'] = (4, 2, config['cond_ae_body']['bottleneck_size'])
+config['cond_ae_body']['decoder_depths'] = [3, 3, 9, 3]
+config['cond_ae_body']['decoder_dims'] = [config['cond_ae_body']['bottleneck_size'], 128, 64, 16]

configs/config_abl_overfit_dress_var_2.py

@@ -0,0 +1,122 @@
+config = dict()
+
+# DON'T FORGET TO RESET SSIM LOSS USAGE
+# DON'T FORGET TO CHECK MODEL LOADING
+# DON'T FORGET TO CHECK IF STORE IN MEMORY
+# DON'T FORGET TO CHECK REG LOSS COEFF
+# DON'T FORGET TO LOAD FINAL UPSCALER UNLESS HAVE LEVEL 512
+# DON'T FORGET TO CHECK AVG GARMENT USED
+
+config['use_skinning'] =              True # and deshaping
+config['remove_global_orientation'] = True
+config['use_curvature'] =             True
+config['use_extra_losses'] =          True # edge, collision, 3d mesh
+config['use_offsets'] =               True
+
+config['use_norm_loss'] =             True
+config['only_curvatures'] =           True
+
+ablation_count = [config['use_skinning'], config['remove_global_orientation'], config['use_curvature'], config['use_extra_losses'], config['use_offsets']].count(False)
+assert ablation_count == 1 or ablation_count == 0
+
+config['gpu'] = '1'
+config['name'] = 'abl_overfit_dress_var_2'
+config['notes'] = 'remove other conditions from input as well, divide values by 1000 before tanh'
+
+config['avg_garment_path'] = 'custom_data/train_dress_30k_no_shape_avg_uv.npy' #custom_data/train_no_shape_avg_uv.npy # train_no_shape_avg_uv  train_no_shape_one_hot_avg_uv
+if config['use_skinning']:
+  assert (config['use_offsets'] and config['avg_garment_path'] != '') or (not config['use_offsets'] and config['avg_garment_path'] == '')
+
+config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
+
+config['batch_size'] = 4
+config['resolution'] = 512
+config['is_skirt'] = True
+config['dataset_in_memory'] = True
+config['is_cape'] = False
+
+config['train_dir'] = '/home/workspace/local-data/train_skirt'
+config['valid_dir'] = '/home/workspace/local-data/train_skirt'
+config['test_dir'] = '/home/workspace/local-data/train_skirt'
+
+config['frames_file'] = 'frames_dress_30k.txt'
+config['frames_file_valid'] = 'frames_dress_30k.txt'
+config['frames_file_test'] = 'frames_dress_30k.txt'
+
+config['static_model'] = 'checkpoints/static_1_skirt'
+config['body_model'] = 'checkpoints/body_1_skirt_2_2'
+config['curvature_model'] = 'checkpoints/curv_1_skirt'
+
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+config['pretrained_model'] = ''
+
+###############    
+config['lr'] = 0.0001#                 0.0001
+###############
+config['starting_balancing_coeff'] = 0.1
+
+config['signature_dynamic'] = { 'id': (), 'frame_nr': (),
+                          'pose_shape': (72 + 11,),
+                          'uv_body_posed': (512, 256, 3),
+                          'mesh_body_posed': (14475, 3,),
+                          'mesh_posed': (14475, 3,),
+                          'uv_curvature': (512, 256, 7),
+                        }
+config['signature_static'] = { 'uv_static': (512, 256, 3),
+                              'mesh_mask': (14475,),
+                              'uv_mask': (512, 256,),
+                              }
+
+if config['use_skinning']:
+  config['signature_dynamic']['uv_unposed'] = (512, 256, 3)
+  config['signature_dynamic']['mesh_unposed'] = (14475, 3,)
+else:
+  config['signature_dynamic']['uv_posed'] = (512, 256, 3)
+
+
+config['upscale_enhancer'] = ''
+config['freeze_encoders'] = True
+config['use_one_hot'] = False
+
+config['cvae'] = dict()
+config['cvae']['is_label'] = False
+config['cvae']['latent_dim'] = config['latent_dim']
+config['cvae']['add_noise'] = True
+config['cvae']['noise_amount'] = 0.001
+
+config['cvae']['kernel_size'] = 7
+config['cvae']['encoder_depths'] = [3, 3, 9, 3]
+config['cvae']['encoder_dims'] = [32, 64, 128, 256]
+config['cvae']['decoder_input'] = (16, 8, 80)
+config['cvae']['decoder_depths'] = [3, 3, 9, 3]
+config['cvae']['decoder_dims'] = [256, 128, 64, 16]
+
+
+config['cond_ae'] = dict()
+config['cond_ae']['is_label'] = False
+config['cond_ae']['upscale_output'] = False
+config['cond_ae']['latent_dim'] = config['latent_dim']
+config['cond_ae']['kernel_size'] = 7
+config['cond_ae']['bottleneck_size'] = 64
+config['cond_ae']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae']['encoder_dims'] = [32, 64, 128, 256, config['cond_ae']['bottleneck_size']]
+config['cond_ae']['decoder_input'] = (8, 4, config['cond_ae']['bottleneck_size'])
+config['cond_ae']['decoder_depths'] = [3, 3, 9, 3, 3]
+config['cond_ae']['decoder_dims'] = [config['cond_ae']['bottleneck_size'], 256, 128, 64, 16]
+
+config['cond_ae_body'] = dict()
+config['cond_ae_body']['latent_dim'] = config['latent_dim']
+config['cond_ae_body']['kernel_size'] = 7
+config['cond_ae_body']['upscale_output'] = True
+config['cond_ae_body']['bottleneck_size'] = 512
+config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3, 3]
+config['cond_ae_body']['encoder_dims'] = [32, 64, 128, 256, 256, config['cond_ae_body']['bottleneck_size']]
+config['cond_ae_body']['decoder_input'] = (4, 2, config['cond_ae_body']['bottleneck_size'])
+config['cond_ae_body']['decoder_depths'] = [3, 3, 9, 3]
+config['cond_ae_body']['decoder_dims'] = [config['cond_ae_body']['bottleneck_size'], 128, 64, 16]

configs/config_abl_overfit_dress_var_3.py

@@ -0,0 +1,122 @@
+config = dict()
+
+# DON'T FORGET TO RESET SSIM LOSS USAGE
+# DON'T FORGET TO CHECK MODEL LOADING
+# DON'T FORGET TO CHECK IF STORE IN MEMORY
+# DON'T FORGET TO CHECK REG LOSS COEFF
+# DON'T FORGET TO LOAD FINAL UPSCALER UNLESS HAVE LEVEL 512
+# DON'T FORGET TO CHECK AVG GARMENT USED
+
+config['use_skinning'] =              True # and deshaping
+config['remove_global_orientation'] = True
+config['use_curvature'] =             True
+config['use_extra_losses'] =          True # edge, collision, 3d mesh
+config['use_offsets'] =               True
+
+config['use_norm_loss'] =             True
+config['only_curvatures'] =           True
+
+ablation_count = [config['use_skinning'], config['remove_global_orientation'], config['use_curvature'], config['use_extra_losses'], config['use_offsets']].count(False)
+assert ablation_count == 1 or ablation_count == 0
+
+config['gpu'] = '1'
+config['name'] = 'abl_overfit_dress_var_3'
+config['notes'] = 'remove other conditions from input as well, divide values by 1000 before tanh, continued from abl_overfit_dress_var_2'
+
+config['avg_garment_path'] = 'custom_data/train_dress_30k_no_shape_avg_uv.npy' #custom_data/train_no_shape_avg_uv.npy # train_no_shape_avg_uv  train_no_shape_one_hot_avg_uv
+if config['use_skinning']:
+  assert (config['use_offsets'] and config['avg_garment_path'] != '') or (not config['use_offsets'] and config['avg_garment_path'] == '')
+
+config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
+
+config['batch_size'] = 4
+config['resolution'] = 512
+config['is_skirt'] = True
+config['dataset_in_memory'] = True
+config['is_cape'] = False
+
+config['train_dir'] = '/home/workspace/local-data/train_skirt'
+config['valid_dir'] = '/home/workspace/local-data/train_skirt'
+config['test_dir'] = '/home/workspace/local-data/train_skirt'
+
+config['frames_file'] = 'frames_dress_30k.txt'
+config['frames_file_valid'] = 'frames_dress_30k.txt'
+config['frames_file_test'] = 'frames_dress_30k.txt'
+
+config['static_model'] = 'checkpoints/static_1_skirt'
+config['body_model'] = 'checkpoints/body_1_skirt_2_2'
+config['curvature_model'] = 'checkpoints/curv_1_skirt'
+
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+config['pretrained_model'] = 'checkpoints/abl_overfit_dress_var_2'
+
+###############    
+config['lr'] = 0.0001#                 0.0001
+###############
+config['starting_balancing_coeff'] = 1.0
+
+config['signature_dynamic'] = { 'id': (), 'frame_nr': (),
+                          'pose_shape': (72 + 11,),
+                          'uv_body_posed': (512, 256, 3),
+                          'mesh_body_posed': (14475, 3,),
+                          'mesh_posed': (14475, 3,),
+                          'uv_curvature': (512, 256, 7),
+                        }
+config['signature_static'] = { 'uv_static': (512, 256, 3),
+                              'mesh_mask': (14475,),
+                              'uv_mask': (512, 256,),
+                              }
+
+if config['use_skinning']:
+  config['signature_dynamic']['uv_unposed'] = (512, 256, 3)
+  config['signature_dynamic']['mesh_unposed'] = (14475, 3,)
+else:
+  config['signature_dynamic']['uv_posed'] = (512, 256, 3)
+
+
+config['upscale_enhancer'] = ''
+config['freeze_encoders'] = True
+config['use_one_hot'] = False
+
+config['cvae'] = dict()
+config['cvae']['is_label'] = False
+config['cvae']['latent_dim'] = config['latent_dim']
+config['cvae']['add_noise'] = True
+config['cvae']['noise_amount'] = 0.001
+
+config['cvae']['kernel_size'] = 7
+config['cvae']['encoder_depths'] = [3, 3, 9, 3]
+config['cvae']['encoder_dims'] = [32, 64, 128, 256]
+config['cvae']['decoder_input'] = (16, 8, 80)
+config['cvae']['decoder_depths'] = [3, 3, 9, 3]
+config['cvae']['decoder_dims'] = [256, 128, 64, 16]
+
+
+config['cond_ae'] = dict()
+config['cond_ae']['is_label'] = False
+config['cond_ae']['upscale_output'] = False
+config['cond_ae']['latent_dim'] = config['latent_dim']
+config['cond_ae']['kernel_size'] = 7
+config['cond_ae']['bottleneck_size'] = 64
+config['cond_ae']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae']['encoder_dims'] = [32, 64, 128, 256, config['cond_ae']['bottleneck_size']]
+config['cond_ae']['decoder_input'] = (8, 4, config['cond_ae']['bottleneck_size'])
+config['cond_ae']['decoder_depths'] = [3, 3, 9, 3, 3]
+config['cond_ae']['decoder_dims'] = [config['cond_ae']['bottleneck_size'], 256, 128, 64, 16]
+
+config['cond_ae_body'] = dict()
+config['cond_ae_body']['latent_dim'] = config['latent_dim']
+config['cond_ae_body']['kernel_size'] = 7
+config['cond_ae_body']['upscale_output'] = True
+config['cond_ae_body']['bottleneck_size'] = 512
+config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3, 3]
+config['cond_ae_body']['encoder_dims'] = [32, 64, 128, 256, 256, config['cond_ae_body']['bottleneck_size']]
+config['cond_ae_body']['decoder_input'] = (4, 2, config['cond_ae_body']['bottleneck_size'])
+config['cond_ae_body']['decoder_depths'] = [3, 3, 9, 3]
+config['cond_ae_body']['decoder_dims'] = [config['cond_ae_body']['bottleneck_size'], 128, 64, 16]

configs/config_abl_overfit_jumpsuit.py

@@ -0,0 +1,122 @@
+config = dict()
+
+# DON'T FORGET TO RESET SSIM LOSS USAGE
+# DON'T FORGET TO CHECK MODEL LOADING
+# DON'T FORGET TO CHECK IF STORE IN MEMORY
+# DON'T FORGET TO CHECK REG LOSS COEFF
+# DON'T FORGET TO LOAD FINAL UPSCALER UNLESS HAVE LEVEL 512
+# DON'T FORGET TO CHECK AVG GARMENT USED
+
+config['use_skinning'] =              True # and deshaping
+config['remove_global_orientation'] = True
+config['use_curvature'] =             True
+config['use_extra_losses'] =          True # edge, collision, 3d mesh
+config['use_offsets'] =               True
+
+config['use_norm_loss'] =             True
+config['only_curvatures'] =           False
+
+ablation_count = [config['use_skinning'], config['remove_global_orientation'], config['use_curvature'], config['use_extra_losses'], config['use_offsets']].count(False)
+assert ablation_count == 1 or ablation_count == 0
+
+config['gpu'] = '5'
+config['name'] = 'abl_overfit_jumpsuit'
+config['notes'] = 'remove other conditions from input as well'
+
+config['avg_garment_path'] = 'custom_data/train_jumpsuit_30k_no_shape_avg_uv.npy' #custom_data/train_no_shape_avg_uv.npy # train_no_shape_avg_uv  train_no_shape_one_hot_avg_uv
+if config['use_skinning']:
+  assert (config['use_offsets'] and config['avg_garment_path'] != '') or (not config['use_offsets'] and config['avg_garment_path'] == '')
+
+config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
+
+config['batch_size'] = 4
+config['resolution'] = 512
+config['is_skirt'] = False
+config['dataset_in_memory'] = True
+config['is_cape'] = False
+
+config['train_dir'] = '/home/workspace/local-data/train'
+config['valid_dir'] = '/home/workspace/local-data/train'
+config['test_dir'] = '/home/workspace/local-data/train'
+
+config['frames_file'] = 'frames_jumpsuit_30k.txt'
+config['frames_file_valid'] = 'frames_jumpsuit_30k.txt'
+config['frames_file_test'] = 'frames_jumpsuit_30k.txt'
+
+config['static_model'] = 'checkpoints/static_1'
+config['body_model'] = 'checkpoints/body_1'
+config['curvature_model'] = 'checkpoints/curv_1'
+
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+config['pretrained_model'] = ''
+
+###############    
+config['lr'] = 0.0001#                 0.0001
+###############
+config['starting_balancing_coeff'] = 0.01
+
+config['signature_dynamic'] = { 'id': (), 'frame_nr': (),
+                          'pose_shape': (72 + 11,),
+                          'uv_body_posed': (512, 256, 3),
+                          'mesh_body_posed': (14475, 3,),
+                          'mesh_posed': (14475, 3,),
+                          'uv_curvature': (512, 256, 7),
+                        }
+config['signature_static'] = { 'uv_static': (512, 256, 3),
+                              'mesh_mask': (14475,),
+                              'uv_mask': (512, 256,),
+                              }
+
+if config['use_skinning']:
+  config['signature_dynamic']['uv_unposed'] = (512, 256, 3)
+  config['signature_dynamic']['mesh_unposed'] = (14475, 3,)
+else:
+  config['signature_dynamic']['uv_posed'] = (512, 256, 3)
+
+
+config['upscale_enhancer'] = ''
+config['freeze_encoders'] = True
+config['use_one_hot'] = False
+
+config['cvae'] = dict()
+config['cvae']['is_label'] = False
+config['cvae']['latent_dim'] = config['latent_dim']
+config['cvae']['add_noise'] = True
+config['cvae']['noise_amount'] = 0.025
+
+config['cvae']['kernel_size'] = 7
+config['cvae']['encoder_depths'] = [3, 3, 9, 3]
+config['cvae']['encoder_dims'] = [32, 64, 128, 256]
+config['cvae']['decoder_input'] = (16, 8, 80)
+config['cvae']['decoder_depths'] = [3, 3, 9, 3]
+config['cvae']['decoder_dims'] = [256, 128, 64, 16]
+
+
+config['cond_ae'] = dict()
+config['cond_ae']['is_label'] = False
+config['cond_ae']['upscale_output'] = False
+config['cond_ae']['latent_dim'] = config['latent_dim']
+config['cond_ae']['kernel_size'] = 7
+config['cond_ae']['bottleneck_size'] = 64
+config['cond_ae']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae']['encoder_dims'] = [32, 64, 128, 256, config['cond_ae']['bottleneck_size']]
+config['cond_ae']['decoder_input'] = (8, 4, config['cond_ae']['bottleneck_size'])
+config['cond_ae']['decoder_depths'] = [3, 3, 9, 3, 3]
+config['cond_ae']['decoder_dims'] = [config['cond_ae']['bottleneck_size'], 256, 128, 64, 16]
+
+config['cond_ae_body'] = dict()
+config['cond_ae_body']['latent_dim'] = config['latent_dim']
+config['cond_ae_body']['kernel_size'] = 7
+config['cond_ae_body']['upscale_output'] = True
+config['cond_ae_body']['bottleneck_size'] = 512
+config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3, 3]
+config['cond_ae_body']['encoder_dims'] = [32, 64, 128, 256, 256, config['cond_ae_body']['bottleneck_size']]
+config['cond_ae_body']['decoder_input'] = (4, 2, config['cond_ae_body']['bottleneck_size'])
+config['cond_ae_body']['decoder_depths'] = [3, 3, 9, 3]
+config['cond_ae_body']['decoder_dims'] = [config['cond_ae_body']['bottleneck_size'], 128, 64, 16]

configs/config_abl_overfit_jumpsuit_2.py

@@ -0,0 +1,122 @@
+config = dict()
+
+# DON'T FORGET TO RESET SSIM LOSS USAGE
+# DON'T FORGET TO CHECK MODEL LOADING
+# DON'T FORGET TO CHECK IF STORE IN MEMORY
+# DON'T FORGET TO CHECK REG LOSS COEFF
+# DON'T FORGET TO LOAD FINAL UPSCALER UNLESS HAVE LEVEL 512
+# DON'T FORGET TO CHECK AVG GARMENT USED
+
+config['use_skinning'] =              True # and deshaping
+config['remove_global_orientation'] = True
+config['use_curvature'] =             True
+config['use_extra_losses'] =          True # edge, collision, 3d mesh
+config['use_offsets'] =               True
+
+config['use_norm_loss'] =             True
+config['only_curvatures'] =           True
+
+ablation_count = [config['use_skinning'], config['remove_global_orientation'], config['use_curvature'], config['use_extra_losses'], config['use_offsets']].count(False)
+assert ablation_count == 1 or ablation_count == 0
+
+config['gpu'] = '5'
+config['name'] = 'abl_overfit_jumpsuit_2'
+config['notes'] = 'remove other conditions from input as well'
+
+config['avg_garment_path'] = 'custom_data/train_jumpsuit_30k_no_shape_avg_uv.npy' #custom_data/train_no_shape_avg_uv.npy # train_no_shape_avg_uv  train_no_shape_one_hot_avg_uv
+if config['use_skinning']:
+  assert (config['use_offsets'] and config['avg_garment_path'] != '') or (not config['use_offsets'] and config['avg_garment_path'] == '')
+
+config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
+
+config['batch_size'] = 4
+config['resolution'] = 512
+config['is_skirt'] = False
+config['dataset_in_memory'] = True
+config['is_cape'] = False
+
+config['train_dir'] = '/home/workspace/local-data/train'
+config['valid_dir'] = '/home/workspace/local-data/train'
+config['test_dir'] = '/home/workspace/local-data/train'
+
+config['frames_file'] = 'frames_jumpsuit_30k.txt'
+config['frames_file_valid'] = 'frames_jumpsuit_30k.txt'
+config['frames_file_test'] = 'frames_jumpsuit_30k.txt'
+
+config['static_model'] = 'checkpoints/static_1'
+config['body_model'] = 'checkpoints/body_1'
+config['curvature_model'] = 'checkpoints/curv_1'
+
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+config['pretrained_model'] = ''
+
+###############    
+config['lr'] = 0.0001#                 0.0001
+###############
+config['starting_balancing_coeff'] = 0.01
+
+config['signature_dynamic'] = { 'id': (), 'frame_nr': (),
+                          'pose_shape': (72 + 11,),
+                          'uv_body_posed': (512, 256, 3),
+                          'mesh_body_posed': (14475, 3,),
+                          'mesh_posed': (14475, 3,),
+                          'uv_curvature': (512, 256, 7),
+                        }
+config['signature_static'] = { 'uv_static': (512, 256, 3),
+                              'mesh_mask': (14475,),
+                              'uv_mask': (512, 256,),
+                              }
+
+if config['use_skinning']:
+  config['signature_dynamic']['uv_unposed'] = (512, 256, 3)
+  config['signature_dynamic']['mesh_unposed'] = (14475, 3,)
+else:
+  config['signature_dynamic']['uv_posed'] = (512, 256, 3)
+
+
+config['upscale_enhancer'] = ''
+config['freeze_encoders'] = True
+config['use_one_hot'] = False
+
+config['cvae'] = dict()
+config['cvae']['is_label'] = False
+config['cvae']['latent_dim'] = config['latent_dim']
+config['cvae']['add_noise'] = True
+config['cvae']['noise_amount'] = 0.025
+
+config['cvae']['kernel_size'] = 7
+config['cvae']['encoder_depths'] = [3, 3, 9, 3]
+config['cvae']['encoder_dims'] = [32, 64, 128, 256]
+config['cvae']['decoder_input'] = (16, 8, 80)
+config['cvae']['decoder_depths'] = [3, 3, 9, 3]
+config['cvae']['decoder_dims'] = [256, 128, 64, 16]
+
+
+config['cond_ae'] = dict()
+config['cond_ae']['is_label'] = False
+config['cond_ae']['upscale_output'] = False
+config['cond_ae']['latent_dim'] = config['latent_dim']
+config['cond_ae']['kernel_size'] = 7
+config['cond_ae']['bottleneck_size'] = 64
+config['cond_ae']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae']['encoder_dims'] = [32, 64, 128, 256, config['cond_ae']['bottleneck_size']]
+config['cond_ae']['decoder_input'] = (8, 4, config['cond_ae']['bottleneck_size'])
+config['cond_ae']['decoder_depths'] = [3, 3, 9, 3, 3]
+config['cond_ae']['decoder_dims'] = [config['cond_ae']['bottleneck_size'], 256, 128, 64, 16]
+
+config['cond_ae_body'] = dict()
+config['cond_ae_body']['latent_dim'] = config['latent_dim']
+config['cond_ae_body']['kernel_size'] = 7
+config['cond_ae_body']['upscale_output'] = True
+config['cond_ae_body']['bottleneck_size'] = 512
+config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3, 3]
+config['cond_ae_body']['encoder_dims'] = [32, 64, 128, 256, 256, config['cond_ae_body']['bottleneck_size']]
+config['cond_ae_body']['decoder_input'] = (4, 2, config['cond_ae_body']['bottleneck_size'])
+config['cond_ae_body']['decoder_depths'] = [3, 3, 9, 3]
+config['cond_ae_body']['decoder_dims'] = [config['cond_ae_body']['bottleneck_size'], 128, 64, 16]

configs/config_ae_body_post.py

@@ -0,0 +1,55 @@
+config = dict()
+
+
+config['batch_size'] = 4
+config['train_dir'] = '/home/workspace/local-data/train'
+config['valid_dir'] = '/home/workspace/local-data/valid'
+config['test_dir'] = '/home/workspace/data/dataset/test_ready'
+config['use_filled'] = False
+config['resolution'] = 512
+config['is_skirt'] = False
+config['dataset_in_memory'] = False
+config['is_cape'] = False
+config['remove_global_orientation'] = True
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+config['cond_ae'] = dict()
+config['cond_ae']['is_label'] = False
+config['cond_ae']['upscale_output'] = False
+config['cond_ae']['latent_dim'] = config['latent_dim']
+config['cond_ae']['kernel_size'] = 7
+config['cond_ae']['bottleneck_size'] = 16
+config['cond_ae']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae']['encoder_dims'] = [16, 32, 64, 128, config['cond_ae']['bottleneck_size']]
+config['cond_ae']['decoder_depths'] = config['cond_ae']['encoder_depths']
+config['cond_ae']['decoder_dims'] = config['cond_ae']['encoder_dims']
+config['cond_ae']['activation'] = 'tanh'
+
+
+config['gpu'] = '2'
+config['lr'] = 0.0001
+config['name'] = 'body_post_2'
+config['notes'] = 'convnext + convnext, | \
+ \
+| \
+    deeper model | \
+| \
+    ' 
+
+config['cond_ae']['upscale_output'] = False
+config['pretrained_model'] = ''
+
+
+config['signature_dynamic'] = { 'id': (), 'frame_nr': (),
+                          'pose_shape': (72 + 11,),
+                          'uv_body_posed': (3, 512, 256,),
+                          'mesh_body_posed': (config['smpl_res'], 3,),
+                        }
+config['signature_static'] = { 
+                         'mesh_mask': (config['smpl_res'],),
+                        }

configs/config_ae_curv_post.py

@@ -0,0 +1,54 @@
+config = dict()
+
+
+config['batch_size'] = 4
+config['train_dir'] = '/home/workspace/local-data/train'
+config['valid_dir'] = '/home/workspace/local-data/valid'
+config['test_dir'] = '/home/workspace/data/dataset/test_ready'
+config['use_filled'] = False
+config['resolution'] = 512
+config['is_skirt'] = False
+config['dataset_in_memory'] = False
+config['is_cape'] = False
+config['remove_global_orientation'] = True
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+config['cond_ae'] = dict()
+config['cond_ae']['is_label'] = True
+config['cond_ae']['upscale_output'] = False
+config['cond_ae']['latent_dim'] = config['latent_dim']
+config['cond_ae']['kernel_size'] = 7
+config['cond_ae']['bottleneck_size'] = 40
+config['cond_ae']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae']['encoder_dims'] = [16, 32, 64, 128, config['cond_ae']['bottleneck_size']]
+config['cond_ae']['decoder_depths'] = config['cond_ae']['encoder_depths']
+config['cond_ae']['decoder_dims'] = config['cond_ae']['encoder_dims']
+config['cond_ae']['activation'] = 'tanh'
+
+
+config['gpu'] = '3'
+config['lr'] = 0.001
+config['name'] = 'curv_post_1'
+config['notes'] = 'convnext + convnext, | \
+ \
+| \
+    deeper model | \
+| \
+    ' 
+
+config['cond_ae']['upscale_output'] = False
+config['pretrained_model'] = ''
+
+
+config['signature_dynamic'] = { 'id': (), 'frame_nr': (),
+                          'uv_curvature': (512, 256, 7),
+                          'mesh_curvature':  (config['smpl_res'])
+                        }
+config['signature_static'] = { 'mesh_mask': (config['smpl_res'],),
+                         'uv_mask': (512, 256,),
+                        }

configs/config_ae_static.py

@@ -1,13 +1,44 @@
-from configs.config_base import config
+config = dict()
 
-config['gpu'] = '4'
+
+config['batch_size'] = 4
+config['train_dir'] = '/home/workspace/local-data/train'
+config['valid_dir'] = '/home/workspace/local-data/valid'
+config['test_dir'] = '/home/workspace/data/dataset/test_ready'
+config['use_filled'] = False
+config['resolution'] = 512
+config['is_skirt'] = False
+config['dataset_in_memory'] = True
+config['is_cape'] = False
+config['remove_global_orientation'] = True
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+config['cond_ae'] = dict()
+config['cond_ae']['is_label'] = False
+config['cond_ae']['upscale_output'] = False
+config['cond_ae']['latent_dim'] = config['latent_dim']
+config['cond_ae']['kernel_size'] = 7
+config['cond_ae']['bottleneck_size'] = 8
+config['cond_ae']['encoder_depths'] = [3, 3, 3, 27, 3]
+config['cond_ae']['encoder_dims'] = [8, 16, 32, 32, config['cond_ae']['bottleneck_size']]
+config['cond_ae']['decoder_depths'] = config['cond_ae']['encoder_depths']
+config['cond_ae']['decoder_dims'] = config['cond_ae']['encoder_dims']
+config['cond_ae']['activation'] = 'linear'
+
+
+config['gpu'] = '5'
 config['lr'] = 0.0005
-config['name'] = 'static_1_5'
+config['name'] = 'static_post_4'
-config['notes'] = 'convnext + convnext, add dense with tanh at the end of reverse convnext, | \
+config['notes'] = 'convnext + convnext, | \
-    flattened latent dim to 128 | \
+ \
-    train on 100 garments | \
+| \
     deeper model | \
-    larger batch size | \
+linear activation| \
     ' 
 
 config['cond_ae']['upscale_output'] = False

configs/config_base.py

@@ -32,7 +32,7 @@ config['cond_ae']['bottleneck_size'] = 64
 config['cond_ae']['encoder_depths'] = [3, 3, 9, 9, 3]
 config['cond_ae']['encoder_dims'] = [32, 64, 128, 256, config['cond_ae']['bottleneck_size']]
 config['cond_ae']['decoder_input'] = (8, 4, config['cond_ae']['bottleneck_size'])
-config['cond_ae']['decoder_depths'] = [3, 3, 9, 3, 3]
+config['cond_ae']['decoder_depths'] = [3, 3, 9, 9, 3]
 config['cond_ae']['decoder_dims'] = [config['cond_ae']['bottleneck_size'], 256, 128, 64, 16]
 
 # config['cond_ae'] = dict()

configs/config_cape_f_03223_SL_2.py

@@ -0,0 +1,122 @@
+config = dict()
+
+# DON'T FORGET TO RESET SSIM LOSS USAGE
+# DON'T FORGET TO CHECK MODEL LOADING
+# DON'T FORGET TO CHECK IF STORE IN MEMORY
+# DON'T FORGET TO CHECK REG LOSS COEFF
+# DON'T FORGET TO LOAD FINAL UPSCALER UNLESS HAVE LEVEL 512
+# DON'T FORGET TO CHECK AVG GARMENT USED
+
+config['use_skinning'] =              True # and deshaping
+config['remove_global_orientation'] = True
+config['use_curvature'] =             True
+config['use_extra_losses'] =          False # edge, collision, 3d mesh
+config['use_offsets'] =               True
+
+config['use_norm_loss'] =             False
+
+ablation_count = [config['use_skinning'], config['remove_global_orientation'], config['use_curvature'], config['use_extra_losses'], config['use_offsets']].count(False)
+assert ablation_count == 1 or ablation_count == 0
+
+config['gpu'] = '0'
+config['name'] = 'cape_f_03223_SL_2'
+config['notes'] = ''
+
+config['avg_garment_path'] = 'custom_data/train_cape_f_03223_SL_2_no_shape_avg_uv.npy' #custom_data/train_no_shape_avg_uv.npy # train_no_shape_avg_uv  train_no_shape_one_hot_avg_uv
+if config['use_skinning']:
+  assert (config['use_offsets'] and config['avg_garment_path'] != '') or (not config['use_offsets'] and config['avg_garment_path'] == '')
+
+config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
+
+config['batch_size'] = 4
+config['resolution'] = 512
+config['dataset_in_memory'] = True
+
+config['train_dir'] = '/home/workspace/local-data/cape_processed/train/female'
+config['valid_dir'] = '/home/workspace/local-data/cape_processed/train/female'
+config['test_dir'] = '/home/workspace/local-data/cape_processed/train/female'
+config['is_skirt'] = False
+config['is_cape'] = True
+
+config['frames_file'] = 'frames_03223_SL_train_2.txt'
+config['frames_file_valid'] = 'frames_03223_SL_test_2.txt'
+config['frames_file_test'] = 'frames_03223_SL_test_2.txt'
+
+config['static_model'] = 'checkpoints/static_1_cape_f'
+config['body_model'] = 'checkpoints/body_1_cape_f'
+config['curvature_model'] = 'checkpoints/curv_1_cape_f'
+
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+config['pretrained_model'] = ''
+
+###############    
+config['lr'] = 0.0001#                 0.0001
+###############
+config['starting_balancing_coeff'] = 0.0001
+config['balancing_activation_limit'] = 0.005
+
+config['signature_dynamic'] = { 'id': (), 'frame_nr': (),
+                          'pose_shape': (72 + 11,),
+                          'uv_body_posed': (512, 256, 3),
+                          'mesh_body_posed': (config['smpl_res'], 3,),
+                          'mesh_posed': (config['smpl_res'], 3,),
+                          'uv_curvature': (512, 256, 7),
+                        }
+config['signature_static'] = { 'uv_static': (512, 256, 3),
+                              'mesh_mask': (config['smpl_res'],),
+                              'uv_mask': (512, 256,),
+                              }
+
+if config['use_skinning']:
+  config['signature_dynamic']['uv_unposed'] = (512, 256, 3)
+  config['signature_dynamic']['mesh_unposed'] = (config['smpl_res'], 3,)
+else:
+  config['signature_dynamic']['uv_posed'] = (512, 256, 3)
+
+
+config['upscale_enhancer'] = ''
+config['freeze_encoders'] = True
+config['use_one_hot'] = False
+
+config['cvae'] = dict()
+config['cvae']['is_label'] = False
+config['cvae']['latent_dim'] = config['latent_dim']
+config['cvae']['add_noise'] = True
+config['cvae']['noise_amount'] = 0.025
+
+config['cvae']['kernel_size'] = 7
+config['cvae']['encoder_depths'] = [3, 3, 9, 3]
+config['cvae']['encoder_dims'] = [32, 64, 128, 256]
+config['cvae']['decoder_input'] = (16, 8, 80)
+config['cvae']['decoder_depths'] = [3, 3, 9, 3]
+config['cvae']['decoder_dims'] = [256, 128, 64, 16]
+
+
+config['cond_ae'] = dict()
+config['cond_ae']['is_label'] = False
+config['cond_ae']['upscale_output'] = False
+config['cond_ae']['latent_dim'] = config['latent_dim']
+config['cond_ae']['kernel_size'] = 7
+config['cond_ae']['bottleneck_size'] = 64
+config['cond_ae']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae']['encoder_dims'] = [32, 64, 128, 256, config['cond_ae']['bottleneck_size']]
+config['cond_ae']['decoder_input'] = (8, 4, config['cond_ae']['bottleneck_size'])
+config['cond_ae']['decoder_depths'] = [3, 3, 9, 3, 3]
+config['cond_ae']['decoder_dims'] = [config['cond_ae']['bottleneck_size'], 256, 128, 64, 16]
+
+config['cond_ae_body'] = dict()
+config['cond_ae_body']['latent_dim'] = config['latent_dim']
+config['cond_ae_body']['kernel_size'] = 7
+config['cond_ae_body']['upscale_output'] = True
+config['cond_ae_body']['bottleneck_size'] = 512
+config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3, 3]
+config['cond_ae_body']['encoder_dims'] = [32, 64, 128, 256, 256, config['cond_ae_body']['bottleneck_size']]
+config['cond_ae_body']['decoder_input'] = (4, 2, config['cond_ae_body']['bottleneck_size'])
+config['cond_ae_body']['decoder_depths'] = [3, 3, 9, 3]
+config['cond_ae_body']['decoder_dims'] = [config['cond_ae_body']['bottleneck_size'], 128, 64, 16]

configs/config_cape_f_03223_SS.py

@@ -0,0 +1,120 @@
+config = dict()
+
+# DON'T FORGET TO RESET SSIM LOSS USAGE
+# DON'T FORGET TO CHECK MODEL LOADING
+# DON'T FORGET TO CHECK IF STORE IN MEMORY
+# DON'T FORGET TO CHECK REG LOSS COEFF
+# DON'T FORGET TO LOAD FINAL UPSCALER UNLESS HAVE LEVEL 512
+# DON'T FORGET TO CHECK AVG GARMENT USED
+
+config['use_skinning'] =              True # and deshaping
+config['remove_global_orientation'] = True
+config['use_curvature'] =             True
+config['use_extra_losses'] =          False # edge, collision, 3d mesh
+config['use_offsets'] =               True
+
+config['use_norm_loss'] =             False
+
+ablation_count = [config['use_skinning'], config['remove_global_orientation'], config['use_curvature'], config['use_extra_losses'], config['use_offsets']].count(False)
+assert ablation_count == 1 or ablation_count == 0
+
+config['gpu'] = '4'
+config['name'] = 'cape_f_03223_SS'
+config['notes'] = ''
+
+config['avg_garment_path'] = 'custom_data/train_cape_f_03223_SS_no_shape_skirt_avg_uv.npy' #custom_data/train_no_shape_avg_uv.npy # train_no_shape_avg_uv  train_no_shape_one_hot_avg_uv
+if config['use_skinning']:
+  assert (config['use_offsets'] and config['avg_garment_path'] != '') or (not config['use_offsets'] and config['avg_garment_path'] == '')
+
+config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
+
+config['batch_size'] = 4
+config['resolution'] = 512
+config['dataset_in_memory'] = True
+
+config['train_dir'] = '/home/workspace/local-data/cape_processed/train/female'
+config['valid_dir'] = '/home/workspace/local-data/cape_processed/test/female'
+config['test_dir'] = '/home/workspace/local-data/cape_processed/test/female'
+config['is_skirt'] = False
+config['is_cape'] = True
+
+config['frames_file'] = 'frames_03223_SS.txt'
+config['frames_file_valid'] = 'frames_03223_SS.txt'
+
+config['static_model'] = 'checkpoints/static_1_cape_f'
+config['body_model'] = 'checkpoints/body_1_cape_f'
+config['curvature_model'] = 'checkpoints/curv_1_cape_f'
+
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+config['pretrained_model'] = ''
+
+###############    
+config['lr'] = 0.0001#                 0.0001
+###############
+config['starting_balancing_coeff'] = 0.0001
+
+config['signature_dynamic'] = { 'id': (), 'frame_nr': (),
+                          'pose_shape': (72 + 11,),
+                          'uv_body_posed': (512, 256, 3),
+                          'mesh_body_posed': (config['smpl_res'], 3,),
+                          'mesh_posed': (config['smpl_res'], 3,),
+                          'uv_curvature': (512, 256, 7),
+                        }
+config['signature_static'] = { 'uv_static': (512, 256, 3),
+                              'mesh_mask': (config['smpl_res'],),
+                              'uv_mask': (512, 256,),
+                              }
+
+if config['use_skinning']:
+  config['signature_dynamic']['uv_unposed'] = (512, 256, 3)
+  config['signature_dynamic']['mesh_unposed'] = (config['smpl_res'], 3,)
+else:
+  config['signature_dynamic']['uv_posed'] = (512, 256, 3)
+
+
+config['upscale_enhancer'] = ''
+config['freeze_encoders'] = True
+config['use_one_hot'] = False
+
+config['cvae'] = dict()
+config['cvae']['is_label'] = False
+config['cvae']['latent_dim'] = config['latent_dim']
+config['cvae']['add_noise'] = True
+config['cvae']['noise_amount'] = 0.025
+
+config['cvae']['kernel_size'] = 7
+config['cvae']['encoder_depths'] = [3, 3, 9, 3]
+config['cvae']['encoder_dims'] = [32, 64, 128, 256]
+config['cvae']['decoder_input'] = (16, 8, 80)
+config['cvae']['decoder_depths'] = [3, 3, 9, 3]
+config['cvae']['decoder_dims'] = [256, 128, 64, 16]
+
+
+config['cond_ae'] = dict()
+config['cond_ae']['is_label'] = False
+config['cond_ae']['upscale_output'] = False
+config['cond_ae']['latent_dim'] = config['latent_dim']
+config['cond_ae']['kernel_size'] = 7
+config['cond_ae']['bottleneck_size'] = 64
+config['cond_ae']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae']['encoder_dims'] = [32, 64, 128, 256, config['cond_ae']['bottleneck_size']]
+config['cond_ae']['decoder_input'] = (8, 4, config['cond_ae']['bottleneck_size'])
+config['cond_ae']['decoder_depths'] = [3, 3, 9, 3, 3]
+config['cond_ae']['decoder_dims'] = [config['cond_ae']['bottleneck_size'], 256, 128, 64, 16]
+
+config['cond_ae_body'] = dict()
+config['cond_ae_body']['latent_dim'] = config['latent_dim']
+config['cond_ae_body']['kernel_size'] = 7
+config['cond_ae_body']['upscale_output'] = True
+config['cond_ae_body']['bottleneck_size'] = 512
+config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3, 3]
+config['cond_ae_body']['encoder_dims'] = [32, 64, 128, 256, 256, config['cond_ae_body']['bottleneck_size']]
+config['cond_ae_body']['decoder_input'] = (4, 2, config['cond_ae_body']['bottleneck_size'])
+config['cond_ae_body']['decoder_depths'] = [3, 3, 9, 3]
+config['cond_ae_body']['decoder_dims'] = [config['cond_ae_body']['bottleneck_size'], 128, 64, 16]

configs/config_cape_f_03223_SS_2.py

@@ -0,0 +1,122 @@
+config = dict()
+
+# DON'T FORGET TO RESET SSIM LOSS USAGE
+# DON'T FORGET TO CHECK MODEL LOADING
+# DON'T FORGET TO CHECK IF STORE IN MEMORY
+# DON'T FORGET TO CHECK REG LOSS COEFF
+# DON'T FORGET TO LOAD FINAL UPSCALER UNLESS HAVE LEVEL 512
+# DON'T FORGET TO CHECK AVG GARMENT USED
+
+config['use_skinning'] =              True # and deshaping
+config['remove_global_orientation'] = True
+config['use_curvature'] =             True
+config['use_extra_losses'] =          False # edge, collision, 3d mesh
+config['use_offsets'] =               True
+
+config['use_norm_loss'] =             False
+
+ablation_count = [config['use_skinning'], config['remove_global_orientation'], config['use_curvature'], config['use_extra_losses'], config['use_offsets']].count(False)
+assert ablation_count == 1 or ablation_count == 0
+
+config['gpu'] = '0'
+config['name'] = 'cape_f_03223_SS_2'
+config['notes'] = ''
+
+config['avg_garment_path'] = 'custom_data/train_cape_f_03223_SS_2_no_shape_avg_uv.npy' #custom_data/train_no_shape_avg_uv.npy # train_no_shape_avg_uv  train_no_shape_one_hot_avg_uv
+if config['use_skinning']:
+  assert (config['use_offsets'] and config['avg_garment_path'] != '') or (not config['use_offsets'] and config['avg_garment_path'] == '')
+
+config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
+
+config['batch_size'] = 4
+config['resolution'] = 512
+config['dataset_in_memory'] = True
+
+config['train_dir'] = '/home/workspace/local-data/cape_processed/train/female'
+config['valid_dir'] = '/home/workspace/local-data/cape_processed/train/female'
+config['test_dir'] = '/home/workspace/local-data/cape_processed/train/female'
+config['is_skirt'] = False
+config['is_cape'] = True
+
+config['frames_file'] = 'frames_03223_SS_train_2.txt'
+config['frames_file_valid'] = 'frames_03223_SS_test_2.txt'
+config['frames_file_test'] = 'frames_03223_SS_test_2.txt'
+
+config['static_model'] = 'checkpoints/static_1_cape_f'
+config['body_model'] = 'checkpoints/body_1_cape_f'
+config['curvature_model'] = 'checkpoints/curv_1_cape_f'
+
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+config['pretrained_model'] = ''
+
+###############    
+config['lr'] = 0.0001#                 0.0001
+###############
+config['starting_balancing_coeff'] = 0.0001
+config['balancing_activation_limit'] = 0.005
+
+config['signature_dynamic'] = { 'id': (), 'frame_nr': (),
+                          'pose_shape': (72 + 11,),
+                          'uv_body_posed': (512, 256, 3),
+                          'mesh_body_posed': (config['smpl_res'], 3,),
+                          'mesh_posed': (config['smpl_res'], 3,),
+                          'uv_curvature': (512, 256, 7),
+                        }
+config['signature_static'] = { 'uv_static': (512, 256, 3),
+                              'mesh_mask': (config['smpl_res'],),
+                              'uv_mask': (512, 256,),
+                              }
+
+if config['use_skinning']:
+  config['signature_dynamic']['uv_unposed'] = (512, 256, 3)
+  config['signature_dynamic']['mesh_unposed'] = (config['smpl_res'], 3,)
+else:
+  config['signature_dynamic']['uv_posed'] = (512, 256, 3)
+
+
+config['upscale_enhancer'] = ''
+config['freeze_encoders'] = True
+config['use_one_hot'] = False
+
+config['cvae'] = dict()
+config['cvae']['is_label'] = False
+config['cvae']['latent_dim'] = config['latent_dim']
+config['cvae']['add_noise'] = True
+config['cvae']['noise_amount'] = 0.025
+
+config['cvae']['kernel_size'] = 7
+config['cvae']['encoder_depths'] = [3, 3, 9, 3]
+config['cvae']['encoder_dims'] = [32, 64, 128, 256]
+config['cvae']['decoder_input'] = (16, 8, 80)
+config['cvae']['decoder_depths'] = [3, 3, 9, 3]
+config['cvae']['decoder_dims'] = [256, 128, 64, 16]
+
+
+config['cond_ae'] = dict()
+config['cond_ae']['is_label'] = False
+config['cond_ae']['upscale_output'] = False
+config['cond_ae']['latent_dim'] = config['latent_dim']
+config['cond_ae']['kernel_size'] = 7
+config['cond_ae']['bottleneck_size'] = 64
+config['cond_ae']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae']['encoder_dims'] = [32, 64, 128, 256, config['cond_ae']['bottleneck_size']]
+config['cond_ae']['decoder_input'] = (8, 4, config['cond_ae']['bottleneck_size'])
+config['cond_ae']['decoder_depths'] = [3, 3, 9, 3, 3]
+config['cond_ae']['decoder_dims'] = [config['cond_ae']['bottleneck_size'], 256, 128, 64, 16]
+
+config['cond_ae_body'] = dict()
+config['cond_ae_body']['latent_dim'] = config['latent_dim']
+config['cond_ae_body']['kernel_size'] = 7
+config['cond_ae_body']['upscale_output'] = True
+config['cond_ae_body']['bottleneck_size'] = 512
+config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3, 3]
+config['cond_ae_body']['encoder_dims'] = [32, 64, 128, 256, 256, config['cond_ae_body']['bottleneck_size']]
+config['cond_ae_body']['decoder_input'] = (4, 2, config['cond_ae_body']['bottleneck_size'])
+config['cond_ae_body']['decoder_depths'] = [3, 3, 9, 3]
+config['cond_ae_body']['decoder_dims'] = [config['cond_ae_body']['bottleneck_size'], 128, 64, 16]

configs/config_cape_m_00032_SL.py

@@ -0,0 +1,121 @@
+config = dict()
+
+# DON'T FORGET TO RESET SSIM LOSS USAGE
+# DON'T FORGET TO CHECK MODEL LOADING
+# DON'T FORGET TO CHECK IF STORE IN MEMORY
+# DON'T FORGET TO CHECK REG LOSS COEFF
+# DON'T FORGET TO LOAD FINAL UPSCALER UNLESS HAVE LEVEL 512
+# DON'T FORGET TO CHECK AVG GARMENT USED
+
+config['use_skinning'] =              True # and deshaping
+config['remove_global_orientation'] = True
+config['use_curvature'] =             True
+config['use_extra_losses'] =          False # edge, collision, 3d mesh
+config['use_offsets'] =               True
+
+config['use_norm_loss'] =             False
+
+ablation_count = [config['use_skinning'], config['remove_global_orientation'], config['use_curvature'], config['use_extra_losses'], config['use_offsets']].count(False)
+assert ablation_count == 1 or ablation_count == 0
+
+config['gpu'] = '1'
+config['name'] = 'cape_m_00032_SL'
+config['notes'] = ''
+
+config['avg_garment_path'] = 'custom_data/train_cape_m_00032_SL_no_shape_avg_uv.npy' #custom_data/train_no_shape_avg_uv.npy # train_no_shape_avg_uv  train_no_shape_one_hot_avg_uv
+if config['use_skinning']:
+  assert (config['use_offsets'] and config['avg_garment_path'] != '') or (not config['use_offsets'] and config['avg_garment_path'] == '')
+
+config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
+
+config['batch_size'] = 4
+config['resolution'] = 512
+config['dataset_in_memory'] = True
+
+config['train_dir'] = '/home/workspace/local-data/cape_processed/train/male'
+config['valid_dir'] = '/home/workspace/local-data/cape_processed/train/male'
+config['test_dir'] = '/home/workspace/local-data/cape_processed/train/male'
+config['is_skirt'] = False
+config['is_cape'] = True
+
+config['frames_file'] = 'frames_00032_SL_train.txt'
+config['frames_file_valid'] = 'frames_00032_SL_test.txt'
+config['frames_file_test'] = 'frames_00032_SL_test.txt'
+
+config['static_model'] = 'checkpoints/static_1_cape_m'
+config['body_model'] = 'checkpoints/body_1_cape_m'
+config['curvature_model'] = 'checkpoints/curv_1_cape_m'
+
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+config['pretrained_model'] = ''
+
+###############    
+config['lr'] = 0.0001#                 0.0001
+###############
+config['starting_balancing_coeff'] = 0.0001
+
+config['signature_dynamic'] = { 'id': (), 'frame_nr': (),
+                          'pose_shape': (72 + 11,),
+                          'uv_body_posed': (512, 256, 3),
+                          'mesh_body_posed': (config['smpl_res'], 3,),
+                          'mesh_posed': (config['smpl_res'], 3,),
+                          'uv_curvature': (512, 256, 7),
+                        }
+config['signature_static'] = { 'uv_static': (512, 256, 3),
+                              'mesh_mask': (config['smpl_res'],),
+                              'uv_mask': (512, 256,),
+                              }
+
+if config['use_skinning']:
+  config['signature_dynamic']['uv_unposed'] = (512, 256, 3)
+  config['signature_dynamic']['mesh_unposed'] = (config['smpl_res'], 3,)
+else:
+  config['signature_dynamic']['uv_posed'] = (512, 256, 3)
+
+
+config['upscale_enhancer'] = ''
+config['freeze_encoders'] = True
+config['use_one_hot'] = False
+
+config['cvae'] = dict()
+config['cvae']['is_label'] = False
+config['cvae']['latent_dim'] = config['latent_dim']
+config['cvae']['add_noise'] = True
+config['cvae']['noise_amount'] = 0.025
+
+config['cvae']['kernel_size'] = 7
+config['cvae']['encoder_depths'] = [3, 3, 9, 3]
+config['cvae']['encoder_dims'] = [32, 64, 128, 256]
+config['cvae']['decoder_input'] = (16, 8, 80)
+config['cvae']['decoder_depths'] = [3, 3, 9, 3]
+config['cvae']['decoder_dims'] = [256, 128, 64, 16]
+
+
+config['cond_ae'] = dict()
+config['cond_ae']['is_label'] = False
+config['cond_ae']['upscale_output'] = False
+config['cond_ae']['latent_dim'] = config['latent_dim']
+config['cond_ae']['kernel_size'] = 7
+config['cond_ae']['bottleneck_size'] = 64
+config['cond_ae']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae']['encoder_dims'] = [32, 64, 128, 256, config['cond_ae']['bottleneck_size']]
+config['cond_ae']['decoder_input'] = (8, 4, config['cond_ae']['bottleneck_size'])
+config['cond_ae']['decoder_depths'] = [3, 3, 9, 3, 3]
+config['cond_ae']['decoder_dims'] = [config['cond_ae']['bottleneck_size'], 256, 128, 64, 16]
+
+config['cond_ae_body'] = dict()
+config['cond_ae_body']['latent_dim'] = config['latent_dim']
+config['cond_ae_body']['kernel_size'] = 7
+config['cond_ae_body']['upscale_output'] = True
+config['cond_ae_body']['bottleneck_size'] = 512
+config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3, 3]
+config['cond_ae_body']['encoder_dims'] = [32, 64, 128, 256, 256, config['cond_ae_body']['bottleneck_size']]
+config['cond_ae_body']['decoder_input'] = (4, 2, config['cond_ae_body']['bottleneck_size'])
+config['cond_ae_body']['decoder_depths'] = [3, 3, 9, 3]
+config['cond_ae_body']['decoder_dims'] = [config['cond_ae_body']['bottleneck_size'], 128, 64, 16]

configs/config_cape_m_00032_SL_2.py

@@ -0,0 +1,122 @@
+config = dict()
+
+# DON'T FORGET TO RESET SSIM LOSS USAGE
+# DON'T FORGET TO CHECK MODEL LOADING
+# DON'T FORGET TO CHECK IF STORE IN MEMORY
+# DON'T FORGET TO CHECK REG LOSS COEFF
+# DON'T FORGET TO LOAD FINAL UPSCALER UNLESS HAVE LEVEL 512
+# DON'T FORGET TO CHECK AVG GARMENT USED
+
+config['use_skinning'] =              True # and deshaping
+config['remove_global_orientation'] = True
+config['use_curvature'] =             True
+config['use_extra_losses'] =          False # edge, collision, 3d mesh
+config['use_offsets'] =               True
+
+config['use_norm_loss'] =             False
+
+ablation_count = [config['use_skinning'], config['remove_global_orientation'], config['use_curvature'], config['use_extra_losses'], config['use_offsets']].count(False)
+assert ablation_count == 1 or ablation_count == 0
+
+config['gpu'] = '0'
+config['name'] = 'cape_m_00032_SL_2'
+config['notes'] = ''
+
+config['avg_garment_path'] = 'custom_data/train_cape_m_00032_SL_2_no_shape_avg_uv.npy' #custom_data/train_no_shape_avg_uv.npy # train_no_shape_avg_uv  train_no_shape_one_hot_avg_uv
+if config['use_skinning']:
+  assert (config['use_offsets'] and config['avg_garment_path'] != '') or (not config['use_offsets'] and config['avg_garment_path'] == '')
+
+config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
+
+config['batch_size'] = 4
+config['resolution'] = 512
+config['dataset_in_memory'] = True
+
+config['train_dir'] = '/home/workspace/local-data/cape_processed/train/male'
+config['valid_dir'] = '/home/workspace/local-data/cape_processed/train/male'
+config['test_dir'] = '/home/workspace/local-data/cape_processed/train/male'
+config['is_skirt'] = False
+config['is_cape'] = True
+
+config['frames_file'] = 'frames_00032_SL_train_2.txt'
+config['frames_file_valid'] = 'frames_00032_SL_test_2.txt'
+config['frames_file_test'] = 'frames_00032_SL_test_2.txt'
+
+config['static_model'] = 'checkpoints/static_1_cape_m'
+config['body_model'] = 'checkpoints/body_1_cape_m'
+config['curvature_model'] = 'checkpoints/curv_1_cape_m'
+
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+config['pretrained_model'] = ''
+
+###############    
+config['lr'] = 0.0001#                 0.0001
+###############
+config['starting_balancing_coeff'] = 0.0001
+config['balancing_activation_limit'] = 0.005
+
+config['signature_dynamic'] = { 'id': (), 'frame_nr': (),
+                          'pose_shape': (72 + 11,),
+                          'uv_body_posed': (512, 256, 3),
+                          'mesh_body_posed': (config['smpl_res'], 3,),
+                          'mesh_posed': (config['smpl_res'], 3,),
+                          'uv_curvature': (512, 256, 7),
+                        }
+config['signature_static'] = { 'uv_static': (512, 256, 3),
+                              'mesh_mask': (config['smpl_res'],),
+                              'uv_mask': (512, 256,),
+                              }
+
+if config['use_skinning']:
+  config['signature_dynamic']['uv_unposed'] = (512, 256, 3)
+  config['signature_dynamic']['mesh_unposed'] = (config['smpl_res'], 3,)
+else:
+  config['signature_dynamic']['uv_posed'] = (512, 256, 3)
+
+
+config['upscale_enhancer'] = ''
+config['freeze_encoders'] = True
+config['use_one_hot'] = False
+
+config['cvae'] = dict()
+config['cvae']['is_label'] = False
+config['cvae']['latent_dim'] = config['latent_dim']
+config['cvae']['add_noise'] = True
+config['cvae']['noise_amount'] = 0.025
+
+config['cvae']['kernel_size'] = 7
+config['cvae']['encoder_depths'] = [3, 3, 9, 3]
+config['cvae']['encoder_dims'] = [32, 64, 128, 256]
+config['cvae']['decoder_input'] = (16, 8, 80)
+config['cvae']['decoder_depths'] = [3, 3, 9, 3]
+config['cvae']['decoder_dims'] = [256, 128, 64, 16]
+
+
+config['cond_ae'] = dict()
+config['cond_ae']['is_label'] = False
+config['cond_ae']['upscale_output'] = False
+config['cond_ae']['latent_dim'] = config['latent_dim']
+config['cond_ae']['kernel_size'] = 7
+config['cond_ae']['bottleneck_size'] = 64
+config['cond_ae']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae']['encoder_dims'] = [32, 64, 128, 256, config['cond_ae']['bottleneck_size']]
+config['cond_ae']['decoder_input'] = (8, 4, config['cond_ae']['bottleneck_size'])
+config['cond_ae']['decoder_depths'] = [3, 3, 9, 3, 3]
+config['cond_ae']['decoder_dims'] = [config['cond_ae']['bottleneck_size'], 256, 128, 64, 16]
+
+config['cond_ae_body'] = dict()
+config['cond_ae_body']['latent_dim'] = config['latent_dim']
+config['cond_ae_body']['kernel_size'] = 7
+config['cond_ae_body']['upscale_output'] = True
+config['cond_ae_body']['bottleneck_size'] = 512
+config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3, 3]
+config['cond_ae_body']['encoder_dims'] = [32, 64, 128, 256, 256, config['cond_ae_body']['bottleneck_size']]
+config['cond_ae_body']['decoder_input'] = (4, 2, config['cond_ae_body']['bottleneck_size'])
+config['cond_ae_body']['decoder_depths'] = [3, 3, 9, 3]
+config['cond_ae_body']['decoder_dims'] = [config['cond_ae_body']['bottleneck_size'], 128, 64, 16]

configs/config_cvae_curv.py

@@ -1,8 +1,12 @@
 config = dict()
 
 config['gpu'] = '2'
-config['name'] = 'cvae_curv'
+config['name'] = 'cvae_curv_8'
-config['notes'] = ''
+config['notes'] = 'no activation for decoder, softmax at the very end, lower balancing, output only 6 values, concatenate background | \
+  same as 3 but use ce loss | \
+  add ce mesh loss | \
+  higher balancing coeff | \
+  start from curv_5, but higher balancing  '
 
 config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
 
@@ -27,12 +31,13 @@ if config['is_cape']:
 
 config['latent_dim'] = 128
 
-config['pretrained_model'] = ''
+config['pretrained_model'] = 'checkpoints/cvae_curv_5'
+config['pretrained_model_iteration'] = '100000'
 
 ###############    
 config['lr'] = 0.0001#                 0.0001
 ###############
-config['starting_balancing_coeff'] = 0.1
+config['starting_balancing_coeff'] = 0.25
 
 config['signature_dynamic'] = { 'id': (), 'frame_nr': (),
                           'pose_shape': (72 + 11,),

configs/post/config_abl_all.py

@@ -0,0 +1,117 @@
+config = dict()
+
+# DON'T FORGET TO RESET SSIM LOSS USAGE
+# DON'T FORGET TO CHECK MODEL LOADING
+# DON'T FORGET TO CHECK IF STORE IN MEMORY
+# DON'T FORGET TO CHECK REG LOSS COEFF
+# DON'T FORGET TO LOAD FINAL UPSCALER UNLESS HAVE LEVEL 512
+# DON'T FORGET TO CHECK AVG GARMENT USED
+
+config['use_skinning'] =              True # and deshaping
+config['remove_global_orientation'] = True
+config['use_curvature'] =             True
+config['use_extra_losses'] =          True # edge, collision, 3d mesh
+config['use_offsets'] =               True
+
+config['use_norm_loss'] =             True
+
+ablation_count = [config['use_skinning'], config['remove_global_orientation'], config['use_curvature'], config['use_extra_losses'], config['use_offsets']].count(False)
+assert ablation_count == 1 or ablation_count == 0
+
+config['gpu'] = '0'
+config['name'] = 'abl_all'
+config['notes'] = ''
+
+config['avg_garment_path'] = 'custom_data/train_no_shape_avg_uv.npy' #custom_data/train_no_shape_avg_uv.npy # train_no_shape_avg_uv  train_no_shape_one_hot_avg_uv
+if config['use_skinning']:
+  assert (config['use_offsets'] and config['avg_garment_path'] != '') or (not config['use_offsets'] and config['avg_garment_path'] == '')
+
+config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
+
+config['batch_size'] = 4
+config['resolution'] = 512
+config['is_skirt'] = False
+config['dataset_in_memory'] = False
+config['is_cape'] = False
+
+config['train_dir'] = '/home/workspace/local-data/train'
+config['valid_dir'] = '/home/workspace/local-data/valid'
+config['test_dir'] = '/home/workspace/data/dataset/test_ready'
+
+config['static_model'] = 'checkpoints/static_1'
+config['body_model'] = 'checkpoints/body_1'
+config['curvature_model'] = 'checkpoints/curv_1'
+
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+config['pretrained_model'] = ''
+
+###############    
+config['lr'] = 0.0001#                 0.0001
+###############
+config['starting_balancing_coeff'] = 0.0001
+
+config['signature_dynamic'] = { 'id': (), 'frame_nr': (),
+                          'pose_shape': (72 + 11,),
+                          'uv_body_posed': (512, 256, 3),
+                          'mesh_body_posed': (14475, 3,),
+                          'mesh_posed': (14475, 3,),
+                          'uv_curvature': (512, 256, 7),
+                        }
+config['signature_static'] = { 'uv_static': (512, 256, 3),
+                              'mesh_mask': (14475,),
+                              'uv_mask': (512, 256,),
+                              }
+
+if config['use_skinning']:
+  config['signature_dynamic']['uv_unposed'] = (512, 256, 3)
+  config['signature_dynamic']['mesh_unposed'] = (14475, 3,)
+else:
+  config['signature_dynamic']['uv_posed'] = (512, 256, 3)
+
+
+config['upscale_enhancer'] = ''
+config['freeze_encoders'] = True
+config['use_one_hot'] = False
+
+config['cvae'] = dict()
+config['cvae']['is_label'] = False
+config['cvae']['latent_dim'] = config['latent_dim']
+config['cvae']['add_noise'] = True
+config['cvae']['noise_amount'] = 0.025
+
+config['cvae']['kernel_size'] = 7
+config['cvae']['encoder_depths'] = [3, 3, 9, 3]
+config['cvae']['encoder_dims'] = [32, 64, 128, 256]
+config['cvae']['decoder_input'] = (16, 8, 80)
+config['cvae']['decoder_depths'] = [3, 3, 9, 3]
+config['cvae']['decoder_dims'] = [256, 128, 64, 16]
+
+
+config['cond_ae'] = dict()
+config['cond_ae']['is_label'] = False
+config['cond_ae']['upscale_output'] = False
+config['cond_ae']['latent_dim'] = config['latent_dim']
+config['cond_ae']['kernel_size'] = 7
+config['cond_ae']['bottleneck_size'] = 64
+config['cond_ae']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae']['encoder_dims'] = [32, 64, 128, 256, config['cond_ae']['bottleneck_size']]
+config['cond_ae']['decoder_input'] = (8, 4, config['cond_ae']['bottleneck_size'])
+config['cond_ae']['decoder_depths'] = [3, 3, 9, 3, 3]
+config['cond_ae']['decoder_dims'] = [config['cond_ae']['bottleneck_size'], 256, 128, 64, 16]
+
+config['cond_ae_body'] = dict()
+config['cond_ae_body']['latent_dim'] = config['latent_dim']
+config['cond_ae_body']['kernel_size'] = 7
+config['cond_ae_body']['upscale_output'] = True
+config['cond_ae_body']['bottleneck_size'] = 512
+config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3, 3]
+config['cond_ae_body']['encoder_dims'] = [32, 64, 128, 256, 256, config['cond_ae_body']['bottleneck_size']]
+config['cond_ae_body']['decoder_input'] = (4, 2, config['cond_ae_body']['bottleneck_size'])
+config['cond_ae_body']['decoder_depths'] = [3, 3, 9, 3]
+config['cond_ae_body']['decoder_dims'] = [config['cond_ae_body']['bottleneck_size'], 128, 64, 16]

configs/post/config_abl_pyr_32.py

@@ -0,0 +1,184 @@
+config = dict()
+
+config['gpu'] = '2'
+config['name'] = 'abl_pyr_32'
+config['notes'] = 'fixed noise handling'
+config['tags'] = ['pyr_post']
+
+# DON'T FORGET TO RESET SSIM LOSS USAGE
+# DON'T FORGET TO CHECK MODEL LOADING
+# DON'T FORGET TO CHECK IF STORE IN MEMORY
+# DON'T FORGET TO CHECK REG LOSS COEFF
+# DON'T FORGET TO LOAD FINAL UPSCALER UNLESS HAVE LEVEL 512
+# DON'T FORGET TO CHECK IF SKIRT
+# DON'T FORGET TO CHECK IF CAPE
+# DON'T FORGET TO CHECK IF USING DOWNSCALE NONE OR _2
+# DON'T FORGET TO CHECK UPSCALING and ENCODER CHECKPOINTS
+
+config['use_skinning'] =              True # and deshaping
+config['remove_global_orientation'] = True
+config['use_extra_losses'] =          True # edge, collision, 3d mesh
+config['use_offsets'] =               False
+config['use_norm_loss'] =             True
+
+config['pretrained_model'] = ''
+config['pretrained_model_iteration'] = ''
+
+config['pyramid'] = dict()
+config['pyramid']['resolutions'] = [32]
+config['pyramid']['pretrained_pyramid'] = ''
+config['pyramid']['final_upscale_enhancer'] = ''
+
+config['train_dir'] = '/home/workspace/local-data/train'
+config['valid_dir'] = '/home/workspace/local-data/valid'
+config['test_dir'] = '/home/workspace/data/dataset/test_ready'
+config['batch_size'] = 4
+config['is_skirt'] = False
+config['is_cape'] = False
+
+config['dataset_in_memory'] = False
+config['resolution'] = None
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['lr'] = 0.0001
+config['starting_balancing_coeff'] = 0.001
+config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
+config['avg_garment_path'] = ''
+
+config['signature_dynamic'] = { 'id': (), 'frame_nr': (),
+                          'pose_shape': (72 + 11,),
+                          'uv_body_posed': (512, 256, 3),
+                          'mesh_body_posed': (config['smpl_res'], 3,),
+                          'mesh_posed': (config['smpl_res'], 3,),
+                          'mesh_unposed': (config['smpl_res'], 3,),
+                          'uv_unposed': (512, 256, 3),
+                          'uv_curvature': (512, 256, 7),
+                        }
+config['signature_static'] = { 'uv_static': (512, 256, 3),
+                              'mesh_mask': (config['smpl_res'],),
+                              'uv_mask': (512, 256,),
+                              }
+
+config['pyramid']['is_label'] = False
+config['pyramid']['add_noise'] = True
+config['pyramid']['noise_amount'] = 0.001
+
+config['pyramid']['loss_weights'] = [1,1,1,1,1]
+
+for res in config['pyramid']['resolutions']:
+  if res != 512:
+      res_modifier = '_' + str(res)
+      config['signature_dynamic']['uv_body_posed' + res_modifier] = (res, res//2, 3)
+      config['signature_dynamic']['uv_unposed' + res_modifier] = (res, res//2, 3)
+      config['signature_dynamic']['uv_curvature' + res_modifier] = (res, res//2, 7)
+      config['signature_static']['uv_static' + res_modifier] = (res, res//2, 3)
+      config['signature_static']['uv_mask' + res_modifier] = (res, res//2)
+
+config['pyramid']['static_model'] = 'checkpoints/post/static_post_2'
+config['pyramid']['body_model'] = 'checkpoints/post/body_post_2'
+config['pyramid']['curvature_model'] = 'checkpoints/post/curv_post_1'
+config['pyramid']['freeze_encoders'] = True
+
+
+config['pyramid']['levels'] = []
+for res in config['pyramid']['resolutions']:
+  if res == 512:
+    level = dict()
+    level['resolution'] = 512
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 7
+    level['encoder_depths'] = [3, 3, 9, 3]
+    level['encoder_dims'] = [32, 64, 128, 64]
+    level['decoder_input'] = (16, 8, 80)
+    level['decoder_depths'] = [3, 3, 9, 3]
+    level['decoder_dims'] = [256, 128, 64, 16]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = ''
+    config['pyramid']['levels'].append(level)
+  elif res == 256:
+    level = dict()
+    level['resolution'] = 256
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 7
+    level['encoder_depths'] = [9, 9, 9, 9]
+    level['encoder_dims'] = [32, 64, 128, 64]
+    level['decoder_input'] = (8, 4, 128)
+    level['decoder_depths'] = [9, 9, 9, 9]
+    level['decoder_dims'] = [256, 128, 64, 16]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = ''#'checkpoints/up_3_d256_512_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  elif res == 128:
+    level = dict()
+    level['resolution'] = 128
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 5
+    level['encoder_depths'] = [16, 9, 9, 9]
+    level['encoder_dims'] = [64, 128, 128, 64]
+    level['decoder_input'] = (4, 2, 128)
+    level['decoder_depths'] = [16, 9, 9, 9]
+    level['decoder_dims'] = [256, 128, 128, 64]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = ''#'checkpoints/up_3_d128_256_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  elif res == 64:
+    level = dict()
+    level['resolution'] = 64
+    level['kernel_size'] = 5
+    level['encoder_depths'] = [9, 9]
+    level['encoder_dims'] = [64, 64]
+    level['decoder_depths'] = level['encoder_depths']
+    level['decoder_dims'] = level['encoder_dims']
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = 'checkpoints/64_512_2'#checkpoints/up_3_d64_128_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  elif res == 32:
+    level = dict()
+    level['resolution'] = 32
+    level['kernel_size'] = 3
+    level['encoder_depths'] = [16]
+    level['encoder_dims'] = [64]
+    level['decoder_depths'] = level['encoder_depths']
+    level['decoder_dims'] = level['encoder_dims']
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = 'checkpoints/32_512_4'#'checkpoints/up_3_d32_64_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  else:
+    raise NotImplementedError
+
+config['cond_ae_static'] = dict()
+config['cond_ae_static']['is_label'] = False
+config['cond_ae_static']['kernel_size'] = 7
+config['cond_ae_static']['bottleneck_size'] = 8
+config['cond_ae_static']['encoder_depths'] = [3, 3, 3, 9, 3]
+config['cond_ae_static']['encoder_dims'] = [16, 32, 64, 64, config['cond_ae_static']['bottleneck_size']]
+config['cond_ae_static']['decoder_depths'] = config['cond_ae_static']['encoder_depths']
+config['cond_ae_static']['decoder_dims'] = config['cond_ae_static']['encoder_dims']
+config['cond_ae_static']['activation'] = 'tanh'
+
+config['cond_ae_curv'] = dict()
+config['cond_ae_curv']['is_label'] = True
+config['cond_ae_curv']['kernel_size'] = 7
+config['cond_ae_curv']['bottleneck_size'] = 40
+config['cond_ae_curv']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae_curv']['encoder_dims'] = [16, 32, 64, 128, config['cond_ae_curv']['bottleneck_size']]
+config['cond_ae_curv']['decoder_depths'] = config['cond_ae_curv']['encoder_depths']
+config['cond_ae_curv']['decoder_dims'] = config['cond_ae_curv']['encoder_dims']
+config['cond_ae_curv']['activation'] = 'tanh'
+
+config['cond_ae_body'] = dict()
+config['cond_ae_body']['is_label'] = False
+config['cond_ae_body']['kernel_size'] = 7
+config['cond_ae_body']['bottleneck_size'] = 16
+config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae_body']['encoder_dims'] = [16, 32, 64, 128, config['cond_ae_body']['bottleneck_size']]
+config['cond_ae_body']['decoder_depths'] = config['cond_ae_body']['encoder_depths']
+config['cond_ae_body']['decoder_dims'] = config['cond_ae_body']['encoder_dims']
+config['cond_ae_body']['activation'] = 'tanh'

configs/post/config_abl_pyr_32_64.py

@@ -0,0 +1,184 @@
+config = dict()
+
+config['gpu'] = '1'
+config['name'] = 'abl_pyr_32_64'
+config['notes'] = 'fixed noise handling'
+config['tags'] = ['pyr_post']
+
+# DON'T FORGET TO RESET SSIM LOSS USAGE
+# DON'T FORGET TO CHECK MODEL LOADING
+# DON'T FORGET TO CHECK IF STORE IN MEMORY
+# DON'T FORGET TO CHECK REG LOSS COEFF
+# DON'T FORGET TO LOAD FINAL UPSCALER UNLESS HAVE LEVEL 512
+# DON'T FORGET TO CHECK IF SKIRT
+# DON'T FORGET TO CHECK IF CAPE
+# DON'T FORGET TO CHECK IF USING DOWNSCALE NONE OR _2
+# DON'T FORGET TO CHECK UPSCALING and ENCODER CHECKPOINTS
+
+config['use_skinning'] =              True # and deshaping
+config['remove_global_orientation'] = True
+config['use_extra_losses'] =          True # edge, collision, 3d mesh
+config['use_offsets'] =               False
+config['use_norm_loss'] =             True
+
+config['pretrained_model'] = ''
+config['pretrained_model_iteration'] = ''
+
+config['pyramid'] = dict()
+config['pyramid']['resolutions'] = [32, 64]
+config['pyramid']['pretrained_pyramid'] = 'checkpoints/postabl_pyr_32'
+config['pyramid']['final_upscale_enhancer'] = ''
+
+config['train_dir'] = '/home/workspace/local-data/train'
+config['valid_dir'] = '/home/workspace/local-data/valid'
+config['test_dir'] = '/home/workspace/data/dataset/test_ready'
+config['batch_size'] = 4
+config['is_skirt'] = False
+config['is_cape'] = False
+
+config['dataset_in_memory'] = False
+config['resolution'] = None
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['lr'] = 0.0001
+config['starting_balancing_coeff'] = 0.001
+config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
+config['avg_garment_path'] = ''
+
+config['signature_dynamic'] = { 'id': (), 'frame_nr': (),
+                          'pose_shape': (72 + 11,),
+                          'uv_body_posed': (512, 256, 3),
+                          'mesh_body_posed': (config['smpl_res'], 3,),
+                          'mesh_posed': (config['smpl_res'], 3,),
+                          'mesh_unposed': (config['smpl_res'], 3,),
+                          'uv_unposed': (512, 256, 3),
+                          'uv_curvature': (512, 256, 7),
+                        }
+config['signature_static'] = { 'uv_static': (512, 256, 3),
+                              'mesh_mask': (config['smpl_res'],),
+                              'uv_mask': (512, 256,),
+                              }
+
+config['pyramid']['is_label'] = False
+config['pyramid']['add_noise'] = True
+config['pyramid']['noise_amount'] = 0.001
+
+config['pyramid']['loss_weights'] = [1,1,1,1,1]
+
+for res in config['pyramid']['resolutions']:
+  if res != 512:
+      res_modifier = '_' + str(res)
+      config['signature_dynamic']['uv_body_posed' + res_modifier] = (res, res//2, 3)
+      config['signature_dynamic']['uv_unposed' + res_modifier] = (res, res//2, 3)
+      config['signature_dynamic']['uv_curvature' + res_modifier] = (res, res//2, 7)
+      config['signature_static']['uv_static' + res_modifier] = (res, res//2, 3)
+      config['signature_static']['uv_mask' + res_modifier] = (res, res//2)
+
+config['pyramid']['static_model'] = 'checkpoints/post/static_post_2'
+config['pyramid']['body_model'] = 'checkpoints/post/body_post_2'
+config['pyramid']['curvature_model'] = 'checkpoints/post/curv_post_1'
+config['pyramid']['freeze_encoders'] = True
+
+
+config['pyramid']['levels'] = []
+for res in config['pyramid']['resolutions']:
+  if res == 512:
+    level = dict()
+    level['resolution'] = 512
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 7
+    level['encoder_depths'] = [3, 3, 9, 3]
+    level['encoder_dims'] = [32, 64, 128, 64]
+    level['decoder_input'] = (16, 8, 80)
+    level['decoder_depths'] = [3, 3, 9, 3]
+    level['decoder_dims'] = [256, 128, 64, 16]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = ''
+    config['pyramid']['levels'].append(level)
+  elif res == 256:
+    level = dict()
+    level['resolution'] = 256
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 7
+    level['encoder_depths'] = [9, 9, 9, 9]
+    level['encoder_dims'] = [32, 64, 128, 64]
+    level['decoder_input'] = (8, 4, 128)
+    level['decoder_depths'] = [9, 9, 9, 9]
+    level['decoder_dims'] = [256, 128, 64, 16]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = ''#'checkpoints/up_3_d256_512_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  elif res == 128:
+    level = dict()
+    level['resolution'] = 128
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 5
+    level['encoder_depths'] = [16, 9, 9, 9]
+    level['encoder_dims'] = [64, 128, 128, 64]
+    level['decoder_input'] = (4, 2, 128)
+    level['decoder_depths'] = [16, 9, 9, 9]
+    level['decoder_dims'] = [256, 128, 128, 64]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = ''#'checkpoints/up_3_d128_256_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  elif res == 64:
+    level = dict()
+    level['resolution'] = 64
+    level['kernel_size'] = 5
+    level['encoder_depths'] = [9, 9]
+    level['encoder_dims'] = [64, 64]
+    level['decoder_depths'] = level['encoder_depths']
+    level['decoder_dims'] = level['encoder_dims']
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = 'checkpoints/64_512_2'#checkpoints/up_3_d64_128_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  elif res == 32:
+    level = dict()
+    level['resolution'] = 32
+    level['kernel_size'] = 3
+    level['encoder_depths'] = [16]
+    level['encoder_dims'] = [64]
+    level['decoder_depths'] = level['encoder_depths']
+    level['decoder_dims'] = level['encoder_dims']
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = 'checkpoints/32_512_4'#'checkpoints/up_3_d32_64_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  else:
+    raise NotImplementedError
+
+config['cond_ae_static'] = dict()
+config['cond_ae_static']['is_label'] = False
+config['cond_ae_static']['kernel_size'] = 7
+config['cond_ae_static']['bottleneck_size'] = 8
+config['cond_ae_static']['encoder_depths'] = [3, 3, 3, 9, 3]
+config['cond_ae_static']['encoder_dims'] = [16, 32, 64, 64, config['cond_ae_static']['bottleneck_size']]
+config['cond_ae_static']['decoder_depths'] = config['cond_ae_static']['encoder_depths']
+config['cond_ae_static']['decoder_dims'] = config['cond_ae_static']['encoder_dims']
+config['cond_ae_static']['activation'] = 'tanh'
+
+config['cond_ae_curv'] = dict()
+config['cond_ae_curv']['is_label'] = True
+config['cond_ae_curv']['kernel_size'] = 7
+config['cond_ae_curv']['bottleneck_size'] = 40
+config['cond_ae_curv']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae_curv']['encoder_dims'] = [16, 32, 64, 128, config['cond_ae_curv']['bottleneck_size']]
+config['cond_ae_curv']['decoder_depths'] = config['cond_ae_curv']['encoder_depths']
+config['cond_ae_curv']['decoder_dims'] = config['cond_ae_curv']['encoder_dims']
+config['cond_ae_curv']['activation'] = 'tanh'
+
+config['cond_ae_body'] = dict()
+config['cond_ae_body']['is_label'] = False
+config['cond_ae_body']['kernel_size'] = 7
+config['cond_ae_body']['bottleneck_size'] = 16
+config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae_body']['encoder_dims'] = [16, 32, 64, 128, config['cond_ae_body']['bottleneck_size']]
+config['cond_ae_body']['decoder_depths'] = config['cond_ae_body']['encoder_depths']
+config['cond_ae_body']['decoder_dims'] = config['cond_ae_body']['encoder_dims']
+config['cond_ae_body']['activation'] = 'tanh'

configs/post/config_abl_pyr_64.py

@@ -0,0 +1,186 @@
+config = dict()
+
+config['gpu'] = '0'
+config['name'] = 'abl_pyr_64'
+config['notes'] = 'fixed noise handling'
+config['tags'] = ['pyr_post']
+
+# DON'T FORGET TO RESET SSIM LOSS USAGE
+# DON'T FORGET TO CHECK MODEL LOADING
+# DON'T FORGET TO CHECK IF STORE IN MEMORY
+# DON'T FORGET TO CHECK REG LOSS COEFF
+# DON'T FORGET TO LOAD FINAL UPSCALER UNLESS HAVE LEVEL 512
+# DON'T FORGET TO CHECK IF SKIRT
+# DON'T FORGET TO CHECK IF CAPE
+# DON'T FORGET TO CHECK IF USING DOWNSCALE NONE OR _2
+# DON'T FORGET TO CHECK UPSCALING and ENCODER CHECKPOINTS
+
+config['use_skinning'] =              True # and deshaping
+config['remove_global_orientation'] = True
+config['use_extra_losses'] =          True # edge, collision, 3d mesh
+config['use_offsets'] =               False
+config['use_norm_loss'] =             True
+
+config['pretrained_model'] = ''
+config['pretrained_model_iteration'] = ''
+
+config['pyramid'] = dict()
+config['pyramid']['resolutions'] = [64]
+config['pyramid']['pretrained_pyramid'] = ''
+config['pyramid']['final_upscale_enhancer'] = ''
+
+config['train_dir'] = '/home/workspace/local-data/train'
+config['valid_dir'] = '/home/workspace/local-data/valid'
+config['test_dir'] = '/home/workspace/data/dataset/test_ready'
+config['batch_size'] = 4
+config['is_skirt'] = False
+config['is_cape'] = False
+
+config['dataset_in_memory'] = False
+config['resolution'] = None
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['lr'] = 0.0001
+config['starting_balancing_coeff'] = 0.001
+config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
+config['avg_garment_path'] = ''
+
+config['signature_dynamic'] = { 'id': (), 'frame_nr': (),
+                          'pose_shape': (72 + 11,),
+                          'uv_body_posed': (512, 256, 3),
+                          'mesh_body_posed': (config['smpl_res'], 3,),
+                          'mesh_posed': (config['smpl_res'], 3,),
+                          'mesh_unposed': (config['smpl_res'], 3,),
+                          'uv_unposed': (512, 256, 3),
+                          'uv_curvature': (512, 256, 7),
+                        }
+config['signature_static'] = { 'uv_static': (512, 256, 3),
+                              'mesh_mask': (config['smpl_res'],),
+                              'uv_mask': (512, 256,),
+                              }
+
+config['pyramid']['is_label'] = False
+config['pyramid']['add_noise'] = True
+config['pyramid']['noise_amount'] = 0.001
+
+config['pyramid']['loss_weights'] = [1,1,1,1,1]
+
+for res in config['pyramid']['resolutions']:
+  if res != 512:
+      res_modifier = '_' + str(res)
+      config['signature_dynamic']['uv_body_posed' + res_modifier] = (res, res//2, 3)
+      config['signature_dynamic']['uv_unposed' + res_modifier] = (res, res//2, 3)
+      config['signature_dynamic']['uv_curvature' + res_modifier] = (res, res//2, 7)
+      config['signature_static']['uv_static' + res_modifier] = (res, res//2, 3)
+      config['signature_static']['uv_mask' + res_modifier] = (res, res//2)
+
+config['pyramid']['static_model'] = 'checkpoints/post/static_post_2'
+config['pyramid']['body_model'] = 'checkpoints/post/body_post_2'
+config['pyramid']['curvature_model'] = 'checkpoints/post/curv_post_1'
+config['pyramid']['freeze_encoders'] = True
+
+
+config['pyramid']['levels'] = []
+for res in config['pyramid']['resolutions']:
+  if res == 512:
+    level = dict()
+    level['resolution'] = 512
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 7
+    level['encoder_depths'] = [3, 3, 9, 3]
+    level['encoder_dims'] = [32, 64, 128, 64]
+    level['decoder_input'] = (16, 8, 80)
+    level['decoder_depths'] = [3, 3, 9, 3]
+    level['decoder_dims'] = [256, 128, 64, 16]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = ''
+    config['pyramid']['levels'].append(level)
+  elif res == 256:
+    level = dict()
+    level['resolution'] = 256
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 7
+    level['encoder_depths'] = [9, 9, 9, 9]
+    level['encoder_dims'] = [32, 64, 128, 64]
+    level['decoder_input'] = (8, 4, 128)
+    level['decoder_depths'] = [9, 9, 9, 9]
+    level['decoder_dims'] = [256, 128, 64, 16]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = ''#'checkpoints/up_3_d256_512_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  elif res == 128:
+    level = dict()
+    level['resolution'] = 128
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 5
+    level['encoder_depths'] = [16, 9, 9, 9]
+    level['encoder_dims'] = [64, 128, 128, 64]
+    level['decoder_input'] = (4, 2, 128)
+    level['decoder_depths'] = [16, 9, 9, 9]
+    level['decoder_dims'] = [256, 128, 128, 64]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = ''#'checkpoints/up_3_d128_256_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  elif res == 64:
+    level = dict()
+    level['resolution'] = 64
+    level['kernel_size'] = 5
+    level['encoder_depths'] = [9, 9]
+    level['encoder_dims'] = [64, 64]
+    level['decoder_depths'] = level['encoder_depths']
+    level['decoder_dims'] = level['encoder_dims']
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = 'checkpoints/64_512_2'#checkpoints/up_3_d64_128_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  # elif res == 32:
+  #   level = dict()
+  #   level['resolution'] = 32
+  #   level['latent_dim'] = config['latent_dim']
+  #   level['kernel_size'] = 3
+  #   level['encoder_depths'] = [9, 9]
+  #   level['encoder_dims'] = [256, 256]
+  #   level['decoder_input'] = (4, 2, 256)
+  #   level['decoder_depths'] = [9, 9]
+  #   level['decoder_dims'] = [256, 256]
+  #   level['avg_garment_path'] = ''
+  #   level['upscale_enhancer'] = ''#'checkpoints/up_3_d32_64_2'
+  #   level['upscale_enhancer_freeze'] = True
+  #   config['pyramid']['levels'].append(level)
+  else:
+    raise NotImplementedError
+
+config['cond_ae_static'] = dict()
+config['cond_ae_static']['is_label'] = False
+config['cond_ae_static']['kernel_size'] = 7
+config['cond_ae_static']['bottleneck_size'] = 8
+config['cond_ae_static']['encoder_depths'] = [3, 3, 3, 9, 3]
+config['cond_ae_static']['encoder_dims'] = [16, 32, 64, 64, config['cond_ae_static']['bottleneck_size']]
+config['cond_ae_static']['decoder_depths'] = config['cond_ae_static']['encoder_depths']
+config['cond_ae_static']['decoder_dims'] = config['cond_ae_static']['encoder_dims']
+config['cond_ae_static']['activation'] = 'tanh'
+
+config['cond_ae_curv'] = dict()
+config['cond_ae_curv']['is_label'] = True
+config['cond_ae_curv']['kernel_size'] = 7
+config['cond_ae_curv']['bottleneck_size'] = 40
+config['cond_ae_curv']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae_curv']['encoder_dims'] = [16, 32, 64, 128, config['cond_ae_curv']['bottleneck_size']]
+config['cond_ae_curv']['decoder_depths'] = config['cond_ae_curv']['encoder_depths']
+config['cond_ae_curv']['decoder_dims'] = config['cond_ae_curv']['encoder_dims']
+config['cond_ae_curv']['activation'] = 'tanh'
+
+config['cond_ae_body'] = dict()
+config['cond_ae_body']['is_label'] = False
+config['cond_ae_body']['kernel_size'] = 7
+config['cond_ae_body']['bottleneck_size'] = 16
+config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae_body']['encoder_dims'] = [16, 32, 64, 128, config['cond_ae_body']['bottleneck_size']]
+config['cond_ae_body']['decoder_depths'] = config['cond_ae_body']['encoder_depths']
+config['cond_ae_body']['decoder_dims'] = config['cond_ae_body']['encoder_dims']
+config['cond_ae_body']['activation'] = 'tanh'

configs/post/config_abl_pyr_64_128.py

@@ -0,0 +1,185 @@
+config = dict()
+
+config['gpu'] = '0'
+config['name'] = 'abl_pyr_64_128'
+config['notes'] = ''
+config['tags'] = ['pyr_post']
+
+# DON'T FORGET TO RESET SSIM LOSS USAGE
+# DON'T FORGET TO CHECK MODEL LOADING
+# DON'T FORGET TO CHECK IF STORE IN MEMORY
+# DON'T FORGET TO CHECK REG LOSS COEFF
+# DON'T FORGET TO LOAD FINAL UPSCALER UNLESS HAVE LEVEL 512
+# DON'T FORGET TO CHECK IF SKIRT
+# DON'T FORGET TO CHECK IF CAPE
+# DON'T FORGET TO CHECK IF USING DOWNSCALE NONE OR _2
+# DON'T FORGET TO CHECK UPSCALING and ENCODER CHECKPOINTS
+
+config['use_skinning'] =              True # and deshaping
+config['remove_global_orientation'] = True
+config['use_extra_losses'] =          True # edge, collision, 3d mesh
+config['use_offsets'] =               False
+config['use_norm_loss'] =             True
+
+config['pretrained_model'] = ''
+config['pretrained_model_iteration'] = ''
+
+config['pyramid'] = dict()
+config['pyramid']['resolutions'] = [64, 128]
+config['pyramid']['pretrained_pyramid'] = 'checkpoints/postabl_pyr_64'
+config['pyramid']['final_upscale_enhancer'] = ''
+
+config['train_dir'] = '/home/workspace/local-data/train'
+config['valid_dir'] = '/home/workspace/local-data/valid'
+config['test_dir'] = '/home/workspace/data/dataset/test_ready'
+config['batch_size'] = 4
+config['is_skirt'] = False
+config['is_cape'] = False
+
+config['dataset_in_memory'] = False
+config['resolution'] = None
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['lr'] = 0.0001
+config['starting_balancing_coeff'] = 0.001
+config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
+config['avg_garment_path'] = ''
+
+config['signature_dynamic'] = { 'id': (), 'frame_nr': (),
+                          'pose_shape': (72 + 11,),
+                          'uv_body_posed': (512, 256, 3),
+                          'mesh_body_posed': (config['smpl_res'], 3,),
+                          'mesh_posed': (config['smpl_res'], 3,),
+                          'mesh_unposed': (config['smpl_res'], 3,),
+                          'uv_unposed': (512, 256, 3),
+                          'uv_curvature': (512, 256, 7),
+                        }
+config['signature_static'] = { 'uv_static': (512, 256, 3),
+                              'mesh_mask': (config['smpl_res'],),
+                              'uv_mask': (512, 256,),
+                              }
+
+config['pyramid']['is_label'] = False
+config['pyramid']['add_noise'] = True
+config['pyramid']['noise_amount'] = 0.001
+
+config['pyramid']['loss_weights'] = [1,1,1,1,1]
+
+for res in config['pyramid']['resolutions']:
+  if res != 512:
+      res_modifier = '_' + str(res)
+      config['signature_dynamic']['uv_body_posed' + res_modifier] = (res, res//2, 3)
+      config['signature_dynamic']['uv_unposed' + res_modifier] = (res, res//2, 3)
+      config['signature_dynamic']['uv_curvature' + res_modifier] = (res, res//2, 7)
+      config['signature_static']['uv_static' + res_modifier] = (res, res//2, 3)
+      config['signature_static']['uv_mask' + res_modifier] = (res, res//2)
+
+config['pyramid']['static_model'] = 'checkpoints/post/static_post_2'
+config['pyramid']['body_model'] = 'checkpoints/post/body_post_2'
+config['pyramid']['curvature_model'] = 'checkpoints/post/curv_post_1'
+config['pyramid']['freeze_encoders'] = True
+
+
+config['pyramid']['levels'] = []
+for res in config['pyramid']['resolutions']:
+  if res == 512:
+    level = dict()
+    level['resolution'] = 512
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 7
+    level['encoder_depths'] = [3, 3, 9, 3]
+    level['encoder_dims'] = [32, 64, 128, 64]
+    level['decoder_input'] = (16, 8, 80)
+    level['decoder_depths'] = [3, 3, 9, 3]
+    level['decoder_dims'] = [256, 128, 64, 16]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = ''
+    config['pyramid']['levels'].append(level)
+  elif res == 256:
+    level = dict()
+    level['resolution'] = 256
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 7
+    level['encoder_depths'] = [9, 9, 9, 9]
+    level['encoder_dims'] = [32, 64, 128, 64]
+    level['decoder_input'] = (8, 4, 128)
+    level['decoder_depths'] = [9, 9, 9, 9]
+    level['decoder_dims'] = [256, 128, 64, 16]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = ''#'checkpoints/up_3_d256_512_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  elif res == 128:
+    level = dict()
+    level['resolution'] = 128
+    level['kernel_size'] = 5
+    level['encoder_depths'] = [9, 9, 9]
+    level['encoder_dims'] = [32, 64, 64]
+    level['decoder_depths'] = level['encoder_depths']
+    level['decoder_dims'] = level['encoder_dims']
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = 'checkpoints/128_512_3'#'checkpoints/up_3_d128_256_2'
+    level['upscale_enhancer_freeze'] = True
+    level['upscale_enhancer_kernel_size'] = 3
+    config['pyramid']['levels'].append(level)
+  elif res == 64:
+    level = dict()
+    level['resolution'] = 64
+    level['kernel_size'] = 5
+    level['encoder_depths'] = [9, 9]
+    level['encoder_dims'] = [64, 64]
+    level['decoder_depths'] = level['encoder_depths']
+    level['decoder_dims'] = level['encoder_dims']
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = 'checkpoints/64_512_2'#checkpoints/up_3_d64_128_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  # elif res == 32:
+  #   level = dict()
+  #   level['resolution'] = 32
+  #   level['latent_dim'] = config['latent_dim']
+  #   level['kernel_size'] = 3
+  #   level['encoder_depths'] = [9, 9]
+  #   level['encoder_dims'] = [256, 256]
+  #   level['decoder_input'] = (4, 2, 256)
+  #   level['decoder_depths'] = [9, 9]
+  #   level['decoder_dims'] = [256, 256]
+  #   level['avg_garment_path'] = ''
+  #   level['upscale_enhancer'] = ''#'checkpoints/up_3_d32_64_2'
+  #   level['upscale_enhancer_freeze'] = True
+  #   config['pyramid']['levels'].append(level)
+  else:
+    raise NotImplementedError
+
+config['cond_ae_static'] = dict()
+config['cond_ae_static']['is_label'] = False
+config['cond_ae_static']['kernel_size'] = 7
+config['cond_ae_static']['bottleneck_size'] = 8
+config['cond_ae_static']['encoder_depths'] = [3, 3, 3, 9, 3]
+config['cond_ae_static']['encoder_dims'] = [16, 32, 64, 64, config['cond_ae_static']['bottleneck_size']]
+config['cond_ae_static']['decoder_depths'] = config['cond_ae_static']['encoder_depths']
+config['cond_ae_static']['decoder_dims'] = config['cond_ae_static']['encoder_dims']
+config['cond_ae_static']['activation'] = 'tanh'
+
+config['cond_ae_curv'] = dict()
+config['cond_ae_curv']['is_label'] = True
+config['cond_ae_curv']['kernel_size'] = 7
+config['cond_ae_curv']['bottleneck_size'] = 40
+config['cond_ae_curv']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae_curv']['encoder_dims'] = [16, 32, 64, 128, config['cond_ae_curv']['bottleneck_size']]
+config['cond_ae_curv']['decoder_depths'] = config['cond_ae_curv']['encoder_depths']
+config['cond_ae_curv']['decoder_dims'] = config['cond_ae_curv']['encoder_dims']
+config['cond_ae_curv']['activation'] = 'tanh'
+
+config['cond_ae_body'] = dict()
+config['cond_ae_body']['is_label'] = False
+config['cond_ae_body']['kernel_size'] = 7
+config['cond_ae_body']['bottleneck_size'] = 16
+config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae_body']['encoder_dims'] = [16, 32, 64, 128, config['cond_ae_body']['bottleneck_size']]
+config['cond_ae_body']['decoder_depths'] = config['cond_ae_body']['encoder_depths']
+config['cond_ae_body']['decoder_dims'] = config['cond_ae_body']['encoder_dims']
+config['cond_ae_body']['activation'] = 'tanh'

configs/post/config_abl_pyr_64_128_2.py

@@ -0,0 +1,185 @@
+config = dict()
+
+config['gpu'] = '3'
+config['name'] = 'abl_pyr_64_128_2'
+config['notes'] = 'use offsets as input for higher levels'
+config['tags'] = ['pyr_post']
+
+# DON'T FORGET TO RESET SSIM LOSS USAGE
+# DON'T FORGET TO CHECK MODEL LOADING
+# DON'T FORGET TO CHECK IF STORE IN MEMORY
+# DON'T FORGET TO CHECK REG LOSS COEFF
+# DON'T FORGET TO LOAD FINAL UPSCALER UNLESS HAVE LEVEL 512
+# DON'T FORGET TO CHECK IF SKIRT
+# DON'T FORGET TO CHECK IF CAPE
+# DON'T FORGET TO CHECK IF USING DOWNSCALE NONE OR _2
+# DON'T FORGET TO CHECK UPSCALING and ENCODER CHECKPOINTS
+
+config['use_skinning'] =              True # and deshaping
+config['remove_global_orientation'] = True
+config['use_extra_losses'] =          True # edge, collision, 3d mesh
+config['use_offsets'] =               False
+config['use_norm_loss'] =             True
+
+config['pretrained_model'] = ''
+config['pretrained_model_iteration'] = ''
+
+config['pyramid'] = dict()
+config['pyramid']['resolutions'] = [64, 128]
+config['pyramid']['pretrained_pyramid'] = 'checkpoints/postabl_pyr_64'
+config['pyramid']['final_upscale_enhancer'] = ''
+
+config['train_dir'] = '/home/workspace/local-data/train'
+config['valid_dir'] = '/home/workspace/local-data/valid'
+config['test_dir'] = '/home/workspace/data/dataset/test_ready'
+config['batch_size'] = 4
+config['is_skirt'] = False
+config['is_cape'] = False
+
+config['dataset_in_memory'] = False
+config['resolution'] = None
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['lr'] = 0.0001
+config['starting_balancing_coeff'] = 0.001
+config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
+config['avg_garment_path'] = ''
+
+config['signature_dynamic'] = { 'id': (), 'frame_nr': (),
+                          'pose_shape': (72 + 11,),
+                          'uv_body_posed': (512, 256, 3),
+                          'mesh_body_posed': (config['smpl_res'], 3,),
+                          'mesh_posed': (config['smpl_res'], 3,),
+                          'mesh_unposed': (config['smpl_res'], 3,),
+                          'uv_unposed': (512, 256, 3),
+                          'uv_curvature': (512, 256, 7),
+                        }
+config['signature_static'] = { 'uv_static': (512, 256, 3),
+                              'mesh_mask': (config['smpl_res'],),
+                              'uv_mask': (512, 256,),
+                              }
+
+config['pyramid']['is_label'] = False
+config['pyramid']['add_noise'] = True
+config['pyramid']['noise_amount'] = 0.001
+
+config['pyramid']['loss_weights'] = [1,1,1,1,1]
+
+for res in config['pyramid']['resolutions']:
+  if res != 512:
+      res_modifier = '_' + str(res)
+      config['signature_dynamic']['uv_body_posed' + res_modifier] = (res, res//2, 3)
+      config['signature_dynamic']['uv_unposed' + res_modifier] = (res, res//2, 3)
+      config['signature_dynamic']['uv_curvature' + res_modifier] = (res, res//2, 7)
+      config['signature_static']['uv_static' + res_modifier] = (res, res//2, 3)
+      config['signature_static']['uv_mask' + res_modifier] = (res, res//2)
+
+config['pyramid']['static_model'] = 'checkpoints/post/static_post_2'
+config['pyramid']['body_model'] = 'checkpoints/post/body_post_2'
+config['pyramid']['curvature_model'] = 'checkpoints/post/curv_post_1'
+config['pyramid']['freeze_encoders'] = True
+
+
+config['pyramid']['levels'] = []
+for res in config['pyramid']['resolutions']:
+  if res == 512:
+    level = dict()
+    level['resolution'] = 512
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 7
+    level['encoder_depths'] = [3, 3, 9, 3]
+    level['encoder_dims'] = [32, 64, 128, 64]
+    level['decoder_input'] = (16, 8, 80)
+    level['decoder_depths'] = [3, 3, 9, 3]
+    level['decoder_dims'] = [256, 128, 64, 16]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = ''
+    config['pyramid']['levels'].append(level)
+  elif res == 256:
+    level = dict()
+    level['resolution'] = 256
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 7
+    level['encoder_depths'] = [9, 9, 9, 9]
+    level['encoder_dims'] = [32, 64, 128, 64]
+    level['decoder_input'] = (8, 4, 128)
+    level['decoder_depths'] = [9, 9, 9, 9]
+    level['decoder_dims'] = [256, 128, 64, 16]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = ''#'checkpoints/up_3_d256_512_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  elif res == 128:
+    level = dict()
+    level['resolution'] = 128
+    level['kernel_size'] = 5
+    level['encoder_depths'] = [9, 9, 9]
+    level['encoder_dims'] = [32, 64, 64]
+    level['decoder_depths'] = level['encoder_depths']
+    level['decoder_dims'] = level['encoder_dims']
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = 'checkpoints/128_512_3'#'checkpoints/up_3_d128_256_2'
+    level['upscale_enhancer_freeze'] = True
+    level['upscale_enhancer_kernel_size'] = 3
+    config['pyramid']['levels'].append(level)
+  elif res == 64:
+    level = dict()
+    level['resolution'] = 64
+    level['kernel_size'] = 5
+    level['encoder_depths'] = [9, 9]
+    level['encoder_dims'] = [64, 64]
+    level['decoder_depths'] = level['encoder_depths']
+    level['decoder_dims'] = level['encoder_dims']
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = 'checkpoints/64_512_2'#checkpoints/up_3_d64_128_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  # elif res == 32:
+  #   level = dict()
+  #   level['resolution'] = 32
+  #   level['latent_dim'] = config['latent_dim']
+  #   level['kernel_size'] = 3
+  #   level['encoder_depths'] = [9, 9]
+  #   level['encoder_dims'] = [256, 256]
+  #   level['decoder_input'] = (4, 2, 256)
+  #   level['decoder_depths'] = [9, 9]
+  #   level['decoder_dims'] = [256, 256]
+  #   level['avg_garment_path'] = ''
+  #   level['upscale_enhancer'] = ''#'checkpoints/up_3_d32_64_2'
+  #   level['upscale_enhancer_freeze'] = True
+  #   config['pyramid']['levels'].append(level)
+  else:
+    raise NotImplementedError
+
+config['cond_ae_static'] = dict()
+config['cond_ae_static']['is_label'] = False
+config['cond_ae_static']['kernel_size'] = 7
+config['cond_ae_static']['bottleneck_size'] = 8
+config['cond_ae_static']['encoder_depths'] = [3, 3, 3, 9, 3]
+config['cond_ae_static']['encoder_dims'] = [16, 32, 64, 64, config['cond_ae_static']['bottleneck_size']]
+config['cond_ae_static']['decoder_depths'] = config['cond_ae_static']['encoder_depths']
+config['cond_ae_static']['decoder_dims'] = config['cond_ae_static']['encoder_dims']
+config['cond_ae_static']['activation'] = 'tanh'
+
+config['cond_ae_curv'] = dict()
+config['cond_ae_curv']['is_label'] = True
+config['cond_ae_curv']['kernel_size'] = 7
+config['cond_ae_curv']['bottleneck_size'] = 40
+config['cond_ae_curv']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae_curv']['encoder_dims'] = [16, 32, 64, 128, config['cond_ae_curv']['bottleneck_size']]
+config['cond_ae_curv']['decoder_depths'] = config['cond_ae_curv']['encoder_depths']
+config['cond_ae_curv']['decoder_dims'] = config['cond_ae_curv']['encoder_dims']
+config['cond_ae_curv']['activation'] = 'tanh'
+
+config['cond_ae_body'] = dict()
+config['cond_ae_body']['is_label'] = False
+config['cond_ae_body']['kernel_size'] = 7
+config['cond_ae_body']['bottleneck_size'] = 16
+config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae_body']['encoder_dims'] = [16, 32, 64, 128, config['cond_ae_body']['bottleneck_size']]
+config['cond_ae_body']['decoder_depths'] = config['cond_ae_body']['encoder_depths']
+config['cond_ae_body']['decoder_dims'] = config['cond_ae_body']['encoder_dims']
+config['cond_ae_body']['activation'] = 'tanh'

configs/post/config_abl_pyr_64_128_3.py

@@ -0,0 +1,185 @@
+config = dict()
+
+config['gpu'] = '3'
+config['name'] = 'abl_pyr_64_128_3'
+config['notes'] = 'use offsets as input for higher levels'
+config['tags'] = ['pyr_post']
+
+# DON'T FORGET TO RESET SSIM LOSS USAGE
+# DON'T FORGET TO CHECK MODEL LOADING
+# DON'T FORGET TO CHECK IF STORE IN MEMORY
+# DON'T FORGET TO CHECK REG LOSS COEFF
+# DON'T FORGET TO LOAD FINAL UPSCALER UNLESS HAVE LEVEL 512
+# DON'T FORGET TO CHECK IF SKIRT
+# DON'T FORGET TO CHECK IF CAPE
+# DON'T FORGET TO CHECK IF USING DOWNSCALE NONE OR _2
+# DON'T FORGET TO CHECK UPSCALING and ENCODER CHECKPOINTS
+
+config['use_skinning'] =              True # and deshaping
+config['remove_global_orientation'] = True
+config['use_extra_losses'] =          True # edge, collision, 3d mesh
+config['use_offsets'] =               False
+config['use_norm_loss'] =             True
+
+config['pretrained_model'] = ''
+config['pretrained_model_iteration'] = ''
+
+config['pyramid'] = dict()
+config['pyramid']['resolutions'] = [64, 128]
+config['pyramid']['pretrained_pyramid'] = 'checkpoints/postabl_pyr_64'
+config['pyramid']['final_upscale_enhancer'] = ''
+
+config['train_dir'] = '/home/workspace/local-data/train'
+config['valid_dir'] = '/home/workspace/local-data/valid'
+config['test_dir'] = '/home/workspace/data/dataset/test_ready'
+config['batch_size'] = 4
+config['is_skirt'] = False
+config['is_cape'] = False
+
+config['dataset_in_memory'] = False
+config['resolution'] = None
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['lr'] = 0.0001
+config['starting_balancing_coeff'] = 0.001
+config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
+config['avg_garment_path'] = ''
+
+config['signature_dynamic'] = { 'id': (), 'frame_nr': (),
+                          'pose_shape': (72 + 11,),
+                          'uv_body_posed': (512, 256, 3),
+                          'mesh_body_posed': (config['smpl_res'], 3,),
+                          'mesh_posed': (config['smpl_res'], 3,),
+                          'mesh_unposed': (config['smpl_res'], 3,),
+                          'uv_unposed': (512, 256, 3),
+                          'uv_curvature': (512, 256, 7),
+                        }
+config['signature_static'] = { 'uv_static': (512, 256, 3),
+                              'mesh_mask': (config['smpl_res'],),
+                              'uv_mask': (512, 256,),
+                              }
+
+config['pyramid']['is_label'] = False
+config['pyramid']['add_noise'] = True
+config['pyramid']['noise_amount'] = 0.001
+
+config['pyramid']['loss_weights'] = [1,1,1,1,1]
+
+for res in config['pyramid']['resolutions']:
+  if res != 512:
+      res_modifier = '_' + str(res)
+      config['signature_dynamic']['uv_body_posed' + res_modifier] = (res, res//2, 3)
+      config['signature_dynamic']['uv_unposed' + res_modifier] = (res, res//2, 3)
+      config['signature_dynamic']['uv_curvature' + res_modifier] = (res, res//2, 7)
+      config['signature_static']['uv_static' + res_modifier] = (res, res//2, 3)
+      config['signature_static']['uv_mask' + res_modifier] = (res, res//2)
+
+config['pyramid']['static_model'] = 'checkpoints/post/static_post_2'
+config['pyramid']['body_model'] = 'checkpoints/post/body_post_2'
+config['pyramid']['curvature_model'] = 'checkpoints/post/curv_post_1'
+config['pyramid']['freeze_encoders'] = True
+
+
+config['pyramid']['levels'] = []
+for res in config['pyramid']['resolutions']:
+  if res == 512:
+    level = dict()
+    level['resolution'] = 512
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 7
+    level['encoder_depths'] = [3, 3, 9, 3]
+    level['encoder_dims'] = [32, 64, 128, 64]
+    level['decoder_input'] = (16, 8, 80)
+    level['decoder_depths'] = [3, 3, 9, 3]
+    level['decoder_dims'] = [256, 128, 64, 16]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = ''
+    config['pyramid']['levels'].append(level)
+  elif res == 256:
+    level = dict()
+    level['resolution'] = 256
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 7
+    level['encoder_depths'] = [9, 9, 9, 9]
+    level['encoder_dims'] = [32, 64, 128, 64]
+    level['decoder_input'] = (8, 4, 128)
+    level['decoder_depths'] = [9, 9, 9, 9]
+    level['decoder_dims'] = [256, 128, 64, 16]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = ''#'checkpoints/up_3_d256_512_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  elif res == 128:
+    level = dict()
+    level['resolution'] = 128
+    level['kernel_size'] = 5
+    level['encoder_depths'] = [9, 9, 9]
+    level['encoder_dims'] = [64, 128, 64]
+    level['decoder_depths'] = level['encoder_depths']
+    level['decoder_dims'] = level['encoder_dims']
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = 'checkpoints/128_512_3'#'checkpoints/up_3_d128_256_2'
+    level['upscale_enhancer_freeze'] = True
+    level['upscale_enhancer_kernel_size'] = 3
+    config['pyramid']['levels'].append(level)
+  elif res == 64:
+    level = dict()
+    level['resolution'] = 64
+    level['kernel_size'] = 5
+    level['encoder_depths'] = [9, 9]
+    level['encoder_dims'] = [64, 64]
+    level['decoder_depths'] = level['encoder_depths']
+    level['decoder_dims'] = level['encoder_dims']
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = 'checkpoints/64_512_2'#checkpoints/up_3_d64_128_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  # elif res == 32:
+  #   level = dict()
+  #   level['resolution'] = 32
+  #   level['latent_dim'] = config['latent_dim']
+  #   level['kernel_size'] = 3
+  #   level['encoder_depths'] = [9, 9]
+  #   level['encoder_dims'] = [256, 256]
+  #   level['decoder_input'] = (4, 2, 256)
+  #   level['decoder_depths'] = [9, 9]
+  #   level['decoder_dims'] = [256, 256]
+  #   level['avg_garment_path'] = ''
+  #   level['upscale_enhancer'] = ''#'checkpoints/up_3_d32_64_2'
+  #   level['upscale_enhancer_freeze'] = True
+  #   config['pyramid']['levels'].append(level)
+  else:
+    raise NotImplementedError
+
+config['cond_ae_static'] = dict()
+config['cond_ae_static']['is_label'] = False
+config['cond_ae_static']['kernel_size'] = 7
+config['cond_ae_static']['bottleneck_size'] = 8
+config['cond_ae_static']['encoder_depths'] = [3, 3, 3, 9, 3]
+config['cond_ae_static']['encoder_dims'] = [16, 32, 64, 64, config['cond_ae_static']['bottleneck_size']]
+config['cond_ae_static']['decoder_depths'] = config['cond_ae_static']['encoder_depths']
+config['cond_ae_static']['decoder_dims'] = config['cond_ae_static']['encoder_dims']
+config['cond_ae_static']['activation'] = 'tanh'
+
+config['cond_ae_curv'] = dict()
+config['cond_ae_curv']['is_label'] = True
+config['cond_ae_curv']['kernel_size'] = 7
+config['cond_ae_curv']['bottleneck_size'] = 40
+config['cond_ae_curv']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae_curv']['encoder_dims'] = [16, 32, 64, 128, config['cond_ae_curv']['bottleneck_size']]
+config['cond_ae_curv']['decoder_depths'] = config['cond_ae_curv']['encoder_depths']
+config['cond_ae_curv']['decoder_dims'] = config['cond_ae_curv']['encoder_dims']
+config['cond_ae_curv']['activation'] = 'tanh'
+
+config['cond_ae_body'] = dict()
+config['cond_ae_body']['is_label'] = False
+config['cond_ae_body']['kernel_size'] = 7
+config['cond_ae_body']['bottleneck_size'] = 16
+config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae_body']['encoder_dims'] = [16, 32, 64, 128, config['cond_ae_body']['bottleneck_size']]
+config['cond_ae_body']['decoder_depths'] = config['cond_ae_body']['encoder_depths']
+config['cond_ae_body']['decoder_dims'] = config['cond_ae_body']['encoder_dims']
+config['cond_ae_body']['activation'] = 'tanh'

configs/post/config_abl_pyr_64_128_4.py

@@ -0,0 +1,185 @@
+config = dict()
+
+config['gpu'] = '0'
+config['name'] = 'abl_pyr_64_128_4'
+config['notes'] = 'use offsets as input for higher levels'
+config['tags'] = ['pyr_post']
+
+# DON'T FORGET TO RESET SSIM LOSS USAGE
+# DON'T FORGET TO CHECK MODEL LOADING
+# DON'T FORGET TO CHECK IF STORE IN MEMORY
+# DON'T FORGET TO CHECK REG LOSS COEFF
+# DON'T FORGET TO LOAD FINAL UPSCALER UNLESS HAVE LEVEL 512
+# DON'T FORGET TO CHECK IF SKIRT
+# DON'T FORGET TO CHECK IF CAPE
+# DON'T FORGET TO CHECK IF USING DOWNSCALE NONE OR _2
+# DON'T FORGET TO CHECK UPSCALING and ENCODER CHECKPOINTS
+
+config['use_skinning'] =              True # and deshaping
+config['remove_global_orientation'] = True
+config['use_extra_losses'] =          True # edge, collision, 3d mesh
+config['use_offsets'] =               False
+config['use_norm_loss'] =             True
+
+config['pretrained_model'] = ''
+config['pretrained_model_iteration'] = ''
+
+config['pyramid'] = dict()
+config['pyramid']['resolutions'] = [64, 128]
+config['pyramid']['pretrained_pyramid'] = 'checkpoints/postabl_pyr_64'
+config['pyramid']['final_upscale_enhancer'] = ''
+
+config['train_dir'] = '/home/workspace/local-data/train'
+config['valid_dir'] = '/home/workspace/local-data/valid'
+config['test_dir'] = '/home/workspace/data/dataset/test_ready'
+config['batch_size'] = 4
+config['is_skirt'] = False
+config['is_cape'] = False
+
+config['dataset_in_memory'] = False
+config['resolution'] = None
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['lr'] = 0.0001
+config['starting_balancing_coeff'] = 0.001
+config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
+config['avg_garment_path'] = ''
+
+config['signature_dynamic'] = { 'id': (), 'frame_nr': (),
+                          'pose_shape': (72 + 11,),
+                          'uv_body_posed': (512, 256, 3),
+                          'mesh_body_posed': (config['smpl_res'], 3,),
+                          'mesh_posed': (config['smpl_res'], 3,),
+                          'mesh_unposed': (config['smpl_res'], 3,),
+                          'uv_unposed': (512, 256, 3),
+                          'uv_curvature': (512, 256, 7),
+                        }
+config['signature_static'] = { 'uv_static': (512, 256, 3),
+                              'mesh_mask': (config['smpl_res'],),
+                              'uv_mask': (512, 256,),
+                              }
+
+config['pyramid']['is_label'] = False
+config['pyramid']['add_noise'] = True
+config['pyramid']['noise_amount'] = 0.001
+
+config['pyramid']['loss_weights'] = [1,1,1,1,1]
+
+for res in config['pyramid']['resolutions']:
+  if res != 512:
+      res_modifier = '_' + str(res)
+      config['signature_dynamic']['uv_body_posed' + res_modifier] = (res, res//2, 3)
+      config['signature_dynamic']['uv_unposed' + res_modifier] = (res, res//2, 3)
+      config['signature_dynamic']['uv_curvature' + res_modifier] = (res, res//2, 7)
+      config['signature_static']['uv_static' + res_modifier] = (res, res//2, 3)
+      config['signature_static']['uv_mask' + res_modifier] = (res, res//2)
+
+config['pyramid']['static_model'] = 'checkpoints/post/static_post_2'
+config['pyramid']['body_model'] = 'checkpoints/post/body_post_2'
+config['pyramid']['curvature_model'] = 'checkpoints/post/curv_post_1'
+config['pyramid']['freeze_encoders'] = True
+
+
+config['pyramid']['levels'] = []
+for res in config['pyramid']['resolutions']:
+  if res == 512:
+    level = dict()
+    level['resolution'] = 512
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 7
+    level['encoder_depths'] = [3, 3, 9, 3]
+    level['encoder_dims'] = [32, 64, 128, 64]
+    level['decoder_input'] = (16, 8, 80)
+    level['decoder_depths'] = [3, 3, 9, 3]
+    level['decoder_dims'] = [256, 128, 64, 16]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = ''
+    config['pyramid']['levels'].append(level)
+  elif res == 256:
+    level = dict()
+    level['resolution'] = 256
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 7
+    level['encoder_depths'] = [9, 9, 9, 9]
+    level['encoder_dims'] = [32, 64, 128, 64]
+    level['decoder_input'] = (8, 4, 128)
+    level['decoder_depths'] = [9, 9, 9, 9]
+    level['decoder_dims'] = [256, 128, 64, 16]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = ''#'checkpoints/up_3_d256_512_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  elif res == 128:
+    level = dict()
+    level['resolution'] = 128
+    level['kernel_size'] = 5
+    level['encoder_depths'] = [9, 9, 9]
+    level['encoder_dims'] = [256, 256, 64]
+    level['decoder_depths'] = level['encoder_depths']
+    level['decoder_dims'] = level['encoder_dims']
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = 'checkpoints/128_512_3'#'checkpoints/up_3_d128_256_2'
+    level['upscale_enhancer_freeze'] = True
+    level['upscale_enhancer_kernel_size'] = 3
+    config['pyramid']['levels'].append(level)
+  elif res == 64:
+    level = dict()
+    level['resolution'] = 64
+    level['kernel_size'] = 5
+    level['encoder_depths'] = [9, 9]
+    level['encoder_dims'] = [64, 64]
+    level['decoder_depths'] = level['encoder_depths']
+    level['decoder_dims'] = level['encoder_dims']
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = 'checkpoints/64_512_2'#checkpoints/up_3_d64_128_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  # elif res == 32:
+  #   level = dict()
+  #   level['resolution'] = 32
+  #   level['latent_dim'] = config['latent_dim']
+  #   level['kernel_size'] = 3
+  #   level['encoder_depths'] = [9, 9]
+  #   level['encoder_dims'] = [256, 256]
+  #   level['decoder_input'] = (4, 2, 256)
+  #   level['decoder_depths'] = [9, 9]
+  #   level['decoder_dims'] = [256, 256]
+  #   level['avg_garment_path'] = ''
+  #   level['upscale_enhancer'] = ''#'checkpoints/up_3_d32_64_2'
+  #   level['upscale_enhancer_freeze'] = True
+  #   config['pyramid']['levels'].append(level)
+  else:
+    raise NotImplementedError
+
+config['cond_ae_static'] = dict()
+config['cond_ae_static']['is_label'] = False
+config['cond_ae_static']['kernel_size'] = 7
+config['cond_ae_static']['bottleneck_size'] = 8
+config['cond_ae_static']['encoder_depths'] = [3, 3, 3, 9, 3]
+config['cond_ae_static']['encoder_dims'] = [16, 32, 64, 64, config['cond_ae_static']['bottleneck_size']]
+config['cond_ae_static']['decoder_depths'] = config['cond_ae_static']['encoder_depths']
+config['cond_ae_static']['decoder_dims'] = config['cond_ae_static']['encoder_dims']
+config['cond_ae_static']['activation'] = 'tanh'
+
+config['cond_ae_curv'] = dict()
+config['cond_ae_curv']['is_label'] = True
+config['cond_ae_curv']['kernel_size'] = 7
+config['cond_ae_curv']['bottleneck_size'] = 40
+config['cond_ae_curv']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae_curv']['encoder_dims'] = [16, 32, 64, 128, config['cond_ae_curv']['bottleneck_size']]
+config['cond_ae_curv']['decoder_depths'] = config['cond_ae_curv']['encoder_depths']
+config['cond_ae_curv']['decoder_dims'] = config['cond_ae_curv']['encoder_dims']
+config['cond_ae_curv']['activation'] = 'tanh'
+
+config['cond_ae_body'] = dict()
+config['cond_ae_body']['is_label'] = False
+config['cond_ae_body']['kernel_size'] = 7
+config['cond_ae_body']['bottleneck_size'] = 16
+config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae_body']['encoder_dims'] = [16, 32, 64, 128, config['cond_ae_body']['bottleneck_size']]
+config['cond_ae_body']['decoder_depths'] = config['cond_ae_body']['encoder_depths']
+config['cond_ae_body']['decoder_dims'] = config['cond_ae_body']['encoder_dims']
+config['cond_ae_body']['activation'] = 'tanh'

configs/post/config_abl_pyr_64_128_5.py

@@ -0,0 +1,185 @@
+config = dict()
+
+config['gpu'] = '0'
+config['name'] = 'abl_pyr_64_128_5'
+config['notes'] = 'use offsets as input for higher levels. use min-max normalization on offsets'
+config['tags'] = ['pyr_post']
+
+# DON'T FORGET TO RESET SSIM LOSS USAGE
+# DON'T FORGET TO CHECK MODEL LOADING
+# DON'T FORGET TO CHECK IF STORE IN MEMORY
+# DON'T FORGET TO CHECK REG LOSS COEFF
+# DON'T FORGET TO LOAD FINAL UPSCALER UNLESS HAVE LEVEL 512
+# DON'T FORGET TO CHECK IF SKIRT
+# DON'T FORGET TO CHECK IF CAPE
+# DON'T FORGET TO CHECK IF USING DOWNSCALE NONE OR _2
+# DON'T FORGET TO CHECK UPSCALING and ENCODER CHECKPOINTS
+
+config['use_skinning'] =              True # and deshaping
+config['remove_global_orientation'] = True
+config['use_extra_losses'] =          True # edge, collision, 3d mesh
+config['use_offsets'] =               False
+config['use_norm_loss'] =             True
+
+config['pretrained_model'] = ''
+config['pretrained_model_iteration'] = ''
+
+config['pyramid'] = dict()
+config['pyramid']['resolutions'] = [64, 128]
+config['pyramid']['pretrained_pyramid'] = 'checkpoints/postabl_pyr_64'
+config['pyramid']['final_upscale_enhancer'] = ''
+
+config['train_dir'] = '/home/workspace/local-data/train'
+config['valid_dir'] = '/home/workspace/local-data/valid'
+config['test_dir'] = '/home/workspace/data/dataset/test_ready'
+config['batch_size'] = 4
+config['is_skirt'] = False
+config['is_cape'] = False
+
+config['dataset_in_memory'] = False
+config['resolution'] = None
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['lr'] = 0.0001
+config['starting_balancing_coeff'] = 0.001
+config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
+config['avg_garment_path'] = ''
+
+config['signature_dynamic'] = { 'id': (), 'frame_nr': (),
+                          'pose_shape': (72 + 11,),
+                          'uv_body_posed': (512, 256, 3),
+                          'mesh_body_posed': (config['smpl_res'], 3,),
+                          'mesh_posed': (config['smpl_res'], 3,),
+                          'mesh_unposed': (config['smpl_res'], 3,),
+                          'uv_unposed': (512, 256, 3),
+                          'uv_curvature': (512, 256, 7),
+                        }
+config['signature_static'] = { 'uv_static': (512, 256, 3),
+                              'mesh_mask': (config['smpl_res'],),
+                              'uv_mask': (512, 256,),
+                              }
+
+config['pyramid']['is_label'] = False
+config['pyramid']['add_noise'] = True
+config['pyramid']['noise_amount'] = 0.001
+
+config['pyramid']['loss_weights'] = [1,1,1,1,1]
+
+for res in config['pyramid']['resolutions']:
+  if res != 512:
+      res_modifier = '_' + str(res)
+      config['signature_dynamic']['uv_body_posed' + res_modifier] = (res, res//2, 3)
+      config['signature_dynamic']['uv_unposed' + res_modifier] = (res, res//2, 3)
+      config['signature_dynamic']['uv_curvature' + res_modifier] = (res, res//2, 7)
+      config['signature_static']['uv_static' + res_modifier] = (res, res//2, 3)
+      config['signature_static']['uv_mask' + res_modifier] = (res, res//2)
+
+config['pyramid']['static_model'] = 'checkpoints/post/static_post_2'
+config['pyramid']['body_model'] = 'checkpoints/post/body_post_2'
+config['pyramid']['curvature_model'] = 'checkpoints/post/curv_post_1'
+config['pyramid']['freeze_encoders'] = True
+
+
+config['pyramid']['levels'] = []
+for res in config['pyramid']['resolutions']:
+  if res == 512:
+    level = dict()
+    level['resolution'] = 512
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 7
+    level['encoder_depths'] = [3, 3, 9, 3]
+    level['encoder_dims'] = [32, 64, 128, 64]
+    level['decoder_input'] = (16, 8, 80)
+    level['decoder_depths'] = [3, 3, 9, 3]
+    level['decoder_dims'] = [256, 128, 64, 16]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = ''
+    config['pyramid']['levels'].append(level)
+  elif res == 256:
+    level = dict()
+    level['resolution'] = 256
+    level['latent_dim'] = config['latent_dim']
+    level['kernel_size'] = 7
+    level['encoder_depths'] = [9, 9, 9, 9]
+    level['encoder_dims'] = [32, 64, 128, 64]
+    level['decoder_input'] = (8, 4, 128)
+    level['decoder_depths'] = [9, 9, 9, 9]
+    level['decoder_dims'] = [256, 128, 64, 16]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = ''#'checkpoints/up_3_d256_512_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  elif res == 128:
+    level = dict()
+    level['resolution'] = 128
+    level['kernel_size'] = 5
+    level['encoder_depths'] = [9, 9, 9]
+    level['encoder_dims'] = [256, 256, 64]
+    level['decoder_depths'] = level['encoder_depths']
+    level['decoder_dims'] = level['encoder_dims']
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = 'checkpoints/128_512_3'#'checkpoints/up_3_d128_256_2'
+    level['upscale_enhancer_freeze'] = True
+    level['upscale_enhancer_kernel_size'] = 3
+    config['pyramid']['levels'].append(level)
+  elif res == 64:
+    level = dict()
+    level['resolution'] = 64
+    level['kernel_size'] = 5
+    level['encoder_depths'] = [9, 9]
+    level['encoder_dims'] = [64, 64]
+    level['decoder_depths'] = level['encoder_depths']
+    level['decoder_dims'] = level['encoder_dims']
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = 'checkpoints/64_512_2'#checkpoints/up_3_d64_128_2'
+    level['upscale_enhancer_freeze'] = True
+    config['pyramid']['levels'].append(level)
+  # elif res == 32:
+  #   level = dict()
+  #   level['resolution'] = 32
+  #   level['latent_dim'] = config['latent_dim']
+  #   level['kernel_size'] = 3
+  #   level['encoder_depths'] = [9, 9]
+  #   level['encoder_dims'] = [256, 256]
+  #   level['decoder_input'] = (4, 2, 256)
+  #   level['decoder_depths'] = [9, 9]
+  #   level['decoder_dims'] = [256, 256]
+  #   level['avg_garment_path'] = ''
+  #   level['upscale_enhancer'] = ''#'checkpoints/up_3_d32_64_2'
+  #   level['upscale_enhancer_freeze'] = True
+  #   config['pyramid']['levels'].append(level)
+  else:
+    raise NotImplementedError
+
+config['cond_ae_static'] = dict()
+config['cond_ae_static']['is_label'] = False
+config['cond_ae_static']['kernel_size'] = 7
+config['cond_ae_static']['bottleneck_size'] = 8
+config['cond_ae_static']['encoder_depths'] = [3, 3, 3, 9, 3]
+config['cond_ae_static']['encoder_dims'] = [16, 32, 64, 64, config['cond_ae_static']['bottleneck_size']]
+config['cond_ae_static']['decoder_depths'] = config['cond_ae_static']['encoder_depths']
+config['cond_ae_static']['decoder_dims'] = config['cond_ae_static']['encoder_dims']
+config['cond_ae_static']['activation'] = 'tanh'
+
+config['cond_ae_curv'] = dict()
+config['cond_ae_curv']['is_label'] = True
+config['cond_ae_curv']['kernel_size'] = 7
+config['cond_ae_curv']['bottleneck_size'] = 40
+config['cond_ae_curv']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae_curv']['encoder_dims'] = [16, 32, 64, 128, config['cond_ae_curv']['bottleneck_size']]
+config['cond_ae_curv']['decoder_depths'] = config['cond_ae_curv']['encoder_depths']
+config['cond_ae_curv']['decoder_dims'] = config['cond_ae_curv']['encoder_dims']
+config['cond_ae_curv']['activation'] = 'tanh'
+
+config['cond_ae_body'] = dict()
+config['cond_ae_body']['is_label'] = False
+config['cond_ae_body']['kernel_size'] = 7
+config['cond_ae_body']['bottleneck_size'] = 16
+config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae_body']['encoder_dims'] = [16, 32, 64, 128, config['cond_ae_body']['bottleneck_size']]
+config['cond_ae_body']['decoder_depths'] = config['cond_ae_body']['encoder_depths']
+config['cond_ae_body']['decoder_dims'] = config['cond_ae_body']['encoder_dims']
+config['cond_ae_body']['activation'] = 'tanh'

convert_to_uv_api.py

@@ -0,0 +1,349 @@
+import os
+import sys
+import glob
+import tqdm
+import math
+import time
+import random
+import shutil
+import collections
+import argparse
+import threading
+import numpy as np
+import scipy as sp
+import trimesh as tm
+import cmapy as cm
+import scipy.io as sio
+import scipy.interpolate as spi
+import scipy.stats as stat
+import cv2
+import plotly.express as ex
+import plotly.graph_objects as go
+from custom_utils.custom_utils import m2uv_array
+
+
+def dilate(img):
+    uv_mask_2 = img != 0
+    m = img.min()
+    # norm = img.max() - img.min()
+    img  = (img - m) * uv_mask_2 #* (np.invert(uv_mask_2) * -1 + uv_mask_2)
+    size = 1
+    shape = cv2.MORPH_CROSS # cv.MORPH_CROSS cv.MORPH_ELLIPSE MORPH_RECT
+    element = cv2.getStructuringElement(shape, (2 * size + 1, 2 * size + 1), (size, size))
+    result = cv2.dilate(img, element, iterations=1)
+    result = (result + m) * uv_mask
+    return result
+
+def erode(img):
+    # m = img.min()
+    # norm = img.max() - img.min()
+    # img  = (img - m) 
+    size = 1
+    shape = cv2.MORPH_CROSS # cv.MORPH_CROSS cv.MORPH_ELLIPSE MORPH_RECT
+    element = cv2.getStructuringElement(shape, (2 * size + 1, 2 * size + 1), (size, size))
+    result = cv2.erode(img, element, iterations=2)
+    # result = result + m
+    return result
+
+def chunks(lst, n):
+    """Yield successive n-sized chunks from lst."""
+    for i in range(0, len(lst), n):
+        yield lst[i:i + n]
+
+
+def convert_to_uv(res, m2uv, uv_map, verts, neighbours=10, 
+                  mask=None, 
+                  method='rbf', 
+                  garment_uv_mask=None, 
+                  uv_map_mask=None, 
+                  uv_counter=None,
+                  uv_map_discreet_unique=None, 
+                  rbf_dest=None, 
+                  uv_mask=None):
+
+    verts_uv = np.zeros((res * 2, res, 3))
+    uv_values = np.zeros((uv_map.shape[0], 3))
+    uv_values[m2uv[:, 0]] = verts
+    uv_values[m2uv[:, 1]] += verts
+    uv_values = uv_values[uv_map_mask, :] #   add garment mask
+    uv_values /= uv_counter[:, np.newaxis]
+    return_garment_uv_mask = False
+
+    if mask is not None:
+        garment_mask = np.zeros((uv_map.shape[0]), dtype=bool)
+        garment_mask[m2uv[:, 0]] = mask
+        garment_mask[m2uv[:, 1]] += mask
+        garment_mask = garment_mask[uv_map_mask] #   add garment mask
+    else:
+        garment_mask = np.sum(uv_values, axis=-1) != 0
+
+    # test
+    verts_uv[uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1], :] = uv_values
+    # verts_uv_save = verts_uv.copy()
+
+    # garment_mask = np.sum(uv_values, axis=-1) != 0
+    unique_uv_values = uv_values[garment_mask, :]
+    uv_map_discreet_unique_unique = uv_map_discreet_unique[garment_mask, :]
+    rbf_dest_current = rbf_dest
+
+    if garment_uv_mask is None and mask is not None:
+        garment_uv_mask = np.zeros((res * 2, res))
+        garment_uv_mask[uv_map_discreet_unique_unique[:, 0], uv_map_discreet_unique_unique[:, 1]] = 255
+        garment_uv_mask = cv2.morphologyEx(garment_uv_mask, cv2.MORPH_CLOSE, cv2.getStructuringElement(cv2.MORPH_RECT, (4 * 2 + 1, 4 * 2 + 1), (4, 4)), iterations=4)
+        garment_uv_mask_border = garment_uv_mask.astype(bool)
+        garment_uv_mask = cv2.dilate(garment_uv_mask, cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2).astype(bool)
+        garment_uv_mask_border = ~garment_uv_mask_border & garment_uv_mask
+        garment_uv_mask = garment_uv_mask * uv_mask[:,:,0]
+        rbf_dest_current = rbf_dest * garment_uv_mask
+        return_garment_uv_mask = True
+
+
+    rbf_dest_current = np.nonzero(rbf_dest_current)
+    rbf_dest_current = np.array(rbf_dest_current).T
+
+    # uv_verts_uniq = np.array(np.nonzero(np.sum(verts_uv, axis=-1))).T
+    if method == 'rbf':
+        interpolator =  spi.RBFInterpolator(y=uv_map_discreet_unique_unique, d=unique_uv_values, neighbors=neighbours, kernel='linear')
+    elif method == 'nearest':
+        interpolator =  spi.NearestNDInterpolator(x=uv_map_discreet_unique_unique, y=unique_uv_values)
+    else:
+        raise NotImplementedError
+
+    final = interpolator(rbf_dest_current)
+
+    result = np.zeros((res * 2, res, 3), np.float32)
+
+    result[rbf_dest_current[:, 0], rbf_dest_current[:, 1], :] = final
+    result[uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1], :] = uv_values
+    if garment_uv_mask is not None:
+        result = result * garment_uv_mask[:,:,np.newaxis]
+
+
+    if return_garment_uv_mask:
+        return result, garment_uv_mask, garment_uv_mask_border
+    else:
+        return result 
+
+
+
+def convert_to_uv_batch(res, m2uv, uv_map, verts, 
+                        neighbours=15, 
+                        mask=None, 
+                        method='rbf', 
+                        garment_uv_mask=None, 
+                        garment_uv_mask_border=None, 
+                        id=None, 
+                        uv_map_mask=None, 
+                        uv_counter=None,
+                        uv_map_discreet_unique=None, 
+                        rbf_dest=None, 
+                        uv_mask=None):
+    batch = len(verts)
+    assert method == 'rbf'
+
+
+    assert mask is not None
+    garment_mask = np.zeros((uv_map.shape[0]), dtype=bool)
+    garment_mask[m2uv[:, 0]] = mask
+    garment_mask[m2uv[:, 1]] += mask
+    garment_mask = garment_mask[uv_map_mask] #   add garment mask
+
+    rbf_dest_current = rbf_dest
+
+    # add back mesh vertices that are not part of the garment
+    uv_map_discreet_unique_unique_inv = uv_map_discreet_unique[np.invert(garment_mask), :]
+    rbf_dest_fill = np.zeros((res * 2, res), dtype=bool)
+    rbf_dest_fill[uv_map_discreet_unique_unique_inv[:, 0], uv_map_discreet_unique_unique_inv[:, 1]] = True
+
+    assert not (garment_uv_mask is None and mask is not None)
+    assert garment_uv_mask is not None
+    rbf_dest_current = (rbf_dest + rbf_dest_fill) * garment_uv_mask + garment_uv_mask_border
+
+    rbf_dest_current = np.nonzero(rbf_dest_current)
+    rbf_dest_current = np.array(rbf_dest_current).T
+
+
+    for i in range(batch):
+        print('processing frame: ' + str(i))
+
+        uv_values = np.zeros((uv_map.shape[0], 3))
+        uv_values[m2uv[:, 0]] = verts[i]
+        uv_values[m2uv[:, 1]] += verts[i]
+        uv_values = uv_values[uv_map_mask, :]
+        uv_values /= uv_counter[:, np.newaxis]
+
+
+        unique_uv_values = uv_values[garment_mask, :]
+        uv_map_discreet_unique_unique = uv_map_discreet_unique[garment_mask, :]
+
+
+        interpolator =  spi.RBFInterpolator(y=uv_map_discreet_unique_unique, d=unique_uv_values[:,:], neighbors=neighbours, kernel='linear')
+        result = np.zeros((res * 2, res, 3), np.float32)
+        result[rbf_dest_current[:, 0], rbf_dest_current[:, 1], :] = interpolator(rbf_dest_current)
+        result[uv_map_discreet_unique_unique[:, 0], uv_map_discreet_unique_unique[:, 1], :] = unique_uv_values[:,:]
+
+        if garment_uv_mask is not None:
+            result *= garment_uv_mask[:,:,np.newaxis]
+            
+        return result 
+
+    assert False    
+
+
+def convert_to_uv_batch_label(res, m2uv, uv_map, values, neighbours=15, mask=None, method='rbf', garment_uv_mask=None, garment_uv_mask_border=None):
+    batch = values.shape[0]
+    values = (values + 1) * mask
+    values_uv = np.zeros((batch, res * 2, res))
+    missing = np.zeros((res * 2, res))
+    uv_values = np.zeros((batch, uv_map.shape[0]))
+    uv_values[:, m2uv[:, 0]] = values
+    uv_values[:, m2uv[:, 1]] = values
+    uv_values = uv_values[:, uv_map_mask]
+    return_garment_uv_mask = False
+    # test
+    # verts_uv_save = verts_uv.copy()
+
+    # unique_uv_values_mask = np.sum(uv_values[0, :, :], axis=-1) != 0
+    # unique_uv_values = uv_values[:, unique_uv_values_mask, :]
+    # uv_map_discreet_unique_unique = uv_map_discreet_unique[unique_uv_values_mask, :]
+
+    if mask is not None:
+        garment_mask = np.zeros((uv_map.shape[0]), dtype=bool)
+        garment_mask[m2uv[:, 0]] = mask
+        garment_mask[m2uv[:, 1]] += mask
+        garment_mask = garment_mask[uv_map_mask] #   add garment mask
+    else:
+        raise Exception
+
+    # uv_values = uv_values * garment_mask
+
+    # test 
+    values_uv[:, uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1]] = uv_values
+    missing[uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1]] = (uv_values[0,:] == 0)
+    
+    unique_uv_values = uv_values[:, garment_mask]
+    uv_map_discreet_unique_unique = uv_map_discreet_unique[garment_mask, :]
+    # rbf_dest_current = rbf_dest.astype(bool) + garment_uv_mask_border * uv_mask[:,:,0]
+
+    if garment_uv_mask is None and mask is not None:
+        garment_uv_mask = np.zeros((res * 2, res))
+        garment_uv_mask[uv_map_discreet_unique_unique[:, 0], uv_map_discreet_unique_unique[:, 1]] = 255
+        garment_uv_mask = cv2.morphologyEx(garment_uv_mask, cv2.MORPH_CLOSE, cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=4)
+        garment_uv_mask = cv2.dilate(garment_uv_mask, cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2).astype(bool)
+        return_garment_uv_mask = True
+
+    rbf_dest_current = rbf_dest * garment_uv_mask + missing * garment_uv_mask
+    rbf_dest_current = np.nonzero(rbf_dest_current)
+    rbf_dest_current = np.array(rbf_dest_current).T
+
+    # uv_verts_uniq = np.array(np.nonzero(np.sum(verts_uv, axis=-1))).T
+    final = []
+    for i in range(batch):
+        print('processing frame: ' + str(i))
+        interpolator =  spi.NearestNDInterpolator(x=uv_map_discreet_unique_unique, y=unique_uv_values[i,:])
+        final.append(interpolator(rbf_dest_current))
+    final = np.asarray(final)
+
+    result = np.zeros((batch, res * 2, res))
+    result[:, uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1]] = uv_values
+    result[:, rbf_dest_current[:, 0], rbf_dest_current[:, 1]] = final
+    if garment_uv_mask is not None:
+        result = result * garment_uv_mask[np.newaxis,:,:]
+
+    if return_garment_uv_mask:
+        return result, garment_uv_mask
+    else:
+        return result 
+
+
+# https://stackoverflow.com/questions/38187286/find-unique-pairs-in-list-of-pairs
+def unique_pairs(a):
+    a.sort(axis=1)
+    b = np.ascontiguousarray(a).view(
+        np.dtype((np.void, a.dtype.itemsize * a.shape[1]))
+    )
+    _, inv, ct = np.unique(b, return_inverse=True, return_counts=True)
+    return ct[inv] == 1
+
+
+def unique_uv_map_mask(uv_map_discreet, uv_res):
+    mask = np.zeros((uv_map_discreet.shape[0]), dtype=bool)
+    temp = np.zeros((uv_res * 2, uv_res))
+    for p in uv_map_discreet:
+        temp[p[0], p[1]] += 1
+    # temp = temp == 1
+    for i, p in enumerate(uv_map_discreet):
+        if temp[p[0], p[1]] == 1:
+            mask[i] = True
+        elif temp[p[0], p[1]] == 2:
+            mask[i] = True
+            temp[p[0], p[1]] = 0
+    return mask
+
+
+def convert_to_uv_api(to_convert, mask, is_skirt=False, uv_res=256):
+    print("skirt" if is_skirt == 1 else "not skirt")
+
+    F = np.load("custom_utils/smpl_faces" + ("_skirt" if is_skirt == 1 else "") + ".npy")
+    uv_map = np.load('custom_utils/stretched_uv_map_' + ("skirt" if is_skirt == 1 else "3") + '.npy')
+    m2uv = np.load('custom_utils/m2uv' + ("_skirt" if is_skirt == 1 else "") + '.npy', allow_pickle=True).item()
+    m2uv = m2uv_array(m2uv)
+    uv_mask = np.load('custom_utils/mask_256x512' + ("_skirt" if is_skirt == 1 else "") + '.npy')
+
+    uv_mask = (cv2.resize(uv_mask * 255.0, (uv_res, uv_res * 2)) / 255).astype(bool)[:, :, np.newaxis]
+    uv_mask = cv2.dilate(uv_mask.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2).astype(bool)[:, :, np.newaxis]
+
+    uv_map_discreet = np.floor(uv_map * uv_res).astype(np.int32)
+    uv_map_mask = unique_uv_map_mask(uv_map_discreet, uv_res)
+
+    uv_counter = np.zeros((uv_map.shape[0]))
+    uv_counter[m2uv[:, 0]] = 1
+    uv_counter[m2uv[:, 1]] += 1
+    uv_counter = uv_counter[uv_map_mask]
+
+    uv_map_discreet_unique = uv_map_discreet[uv_map_mask, :]
+
+    rbf_dest = np.ones((uv_res * 2, uv_res))
+    rbf_dest[uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1]] = 0
+    rbf_dest[np.invert(uv_mask[:,:,0])] = 0
+
+    _, garment_uv_mask, garment_uv_mask_border = convert_to_uv(uv_res, m2uv, uv_map, to_convert, 
+                             neighbours=3, 
+                             mask=mask, 
+                             method='rbf', 
+                             garment_uv_mask=None, 
+                             uv_map_mask=uv_map_mask, 
+                             uv_counter=uv_counter,
+                             uv_map_discreet_unique=uv_map_discreet_unique, 
+                             rbf_dest=rbf_dest, 
+                             uv_mask=uv_mask)
+
+    uv = convert_to_uv_batch(uv_res, m2uv, uv_map, to_convert[np.newaxis,:,:], 
+                             neighbours=3, 
+                             mask=mask, 
+                             method='rbf', 
+                             garment_uv_mask=garment_uv_mask, 
+                             garment_uv_mask_border=garment_uv_mask_border, 
+                             uv_map_mask=uv_map_mask, 
+                             uv_counter=uv_counter,
+                             uv_map_discreet_unique=uv_map_discreet_unique, 
+                             rbf_dest=rbf_dest, 
+                             uv_mask=uv_mask)
+
+    return uv
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-n", "--name", help = "end of range", default='lhunor')
+    args = parser.parse_args()
+
+
+    start = time.perf_counter()
+
+    uv = convert_to_uv_api(np.random.random((14475, 3)) - 0.5, np.ones((14475), bool), is_skirt=False)
+
+    print(time.perf_counter() - start)
+
+    print()
+

convnext/models/convnext_tf.py

@@ -29,7 +29,7 @@ class DropPath(tf.keras.layers.Layer):
         super().__init__()
         self.drop_prob = drop_prob
 
-    def call(self, x, training=None):
+    def call(self, x, training):
         return drop_path(x, self.drop_prob, training)
 
 
@@ -66,19 +66,19 @@ class Block(Layer):
             name=f'{self.prefix}/gamma')
         self.built = True
 
-    def call(self, x):
+    def call(self, x, training):
         input = x
-        x = self.dwconv(x)
+        x = self.dwconv(x, training=training)
         # x = x.permute(0, 2, 3, 1) # (N, C, H, W) -> (N, H, W, C)
-        x = self.norm(x)
+        x = self.norm(x, training=training)
-        x = self.pwconv1(x)
+        x = self.pwconv1(x, training=training)
         x = self.act(x)
-        x = self.pwconv2(x)
+        x = self.pwconv2(x, training=training)
         if self.gamma is not None:
             x = self.gamma * x
         # x = x.permute(0, 3, 1, 2) # (N, H, W, C) -> (N, C, H, W)
 
-        x = input + self.drop_path(x)
+        x = input + self.drop_path(x, training=training)
         return x
 
 
@@ -138,16 +138,16 @@ class ConvNeXt(tf.keras.Model):
             self.norm = None
             self.head = None
 
-    def forward_features(self, x):
+    def forward_features(self, x, training):
         for i in range(len(self.downsample_layers)):
             # x = self.downsample_layers[i](x)
-            x = self.downsample_layers[i][0](x)
+            x = self.downsample_layers[i][0](x, training=training)
-            x = self.downsample_layers[i][1](x)
+            x = self.downsample_layers[i][1](x, training=training)
-            x = self.stages[i](x)
+            x = self.stages[i](x, training=training)
         return x
 
-    def call(self, x):
+    def call(self, x, training):
-        x = self.forward_features(x)
+        x = self.forward_features(x, training=training)
         return x
 
 
@@ -257,11 +257,11 @@ class ConvNeXtReverse(tf.keras.Model):
 
         self.last = Dense(3, activation='tanh')
 
-    def call(self, x):
+    def call(self, x, training):
         for i in reversed(range(len(self.upsample_layers))):
-            x = self.upsample_layers[i](x)
+            x = self.upsample_layers[i](x, training=training)
-            x = self.stages[i](x)
+            x = self.stages[i](x, training=training)
-        x = self.last(x)
+        x = self.last(x, training=training)
         return x
 
 
@@ -366,13 +366,14 @@ class ConvNeXtReverseLabels(tf.keras.Model):
             self.stages.append(stage)
             cur += depths[i]
 
-        self.last = Dense(nr_of_classes, activation='softmax')
+        # self.last = Dense(nr_of_classes, activation='softmax')
+        self.last = Dense(nr_of_classes - 1)
 
-    def call(self, x):
+    def call(self, x, training):
         for i in reversed(range(len(self.downsample_layers))):
-            x = self.downsample_layers[i](x)
+            x = self.downsample_layers[i](x, training=training)
-            x = self.stages[i](x)
+            x = self.stages[i](x, training=training)
-        x = self.last(x)
+        x = self.last(x, training=training)
         return x
 
 

convunext.py

@@ -413,7 +413,7 @@ if __name__ == '__main__':
 
 
     factory = DatasetFactory()
-    signature_dynamic = { 'id': (),
+    signature_dynamic = { 'id': (), 'frame_nr': (),
                           'pose_shape': (72 + 11,),
                         #   'uv_posed': (3, 512, 256,),
                           'uv_body_posed': (3, 512, 256,),

custom_utils/create_uv_map.py

@@ -0,0 +1,596 @@
+import os
+import sys
+import shutil
+import tqdm
+import random
+import math
+import collections
+import tps
+import numpy as np
+import scipy as sp
+import trimesh as tm
+import cmapy as cm
+import scipy.io as sio
+import matplotlib.pyplot as plt
+import plotly.express as ex
+
+
+
+def m2uv_array(m2uv, is_smpl_orig=False):
+    result = np.zeros((14475 if not is_smpl_orig else 3638, 2), dtype=np.int32)
+    for k, v in m2uv.items():
+        if len(v) == 2:
+            result[k] = list(v)
+        elif len(v) == 1:
+            x = v.pop()
+            result[k] = [x, x]
+        else:
+            assert False
+
+    return result
+
+
+def get_body():
+    smpl = SMPL('FeatureLines/superSMPL')
+    pose = np.zeros((1, 72), dtype=np.float32)
+    beta = np.zeros((1, 11), dtype=np.float32)
+    R = smpl.run_rodrigues(pose)
+    T, _, _ = smpl.run(R, beta, True)
+    return np.squeeze(T)
+
+
+def plot_3d(x, y=None):
+    fig = plt.figure(figsize=(15, 15))
+    ax = fig.add_subplot(111, projection='3d')
+    ax.set_xlim3d(-1, 1)
+    ax.set_ylim3d(-1, 1)
+    ax.set_zlim3d(-1, 1)
+    ax.set_xlabel("X")
+    ax.set_ylabel("Y")
+    ax.set_zlabel("Z")
+    ax.scatter(x[:, 0], x[:, 1], x[:, 2], s=5)
+    if y is not None:
+        ax.scatter(y[:, 0], y[:, 1], y[:, 2], s=5)
+    plt.show()
+
+def plot_uv(x, y=None):
+    fig = plt.figure(figsize=(15, 15))
+    ax = fig.add_subplot(111)
+    ax.set_xlim(0, 1)
+    ax.set_ylim(0, 1)
+    ax.scatter(x[:, 0], x[:, 1])
+    if y is not None:
+        ax.scatter(y[:, 0], y[:, 1])
+    plt.show()
+
+""" 
+Mesh to UV map
+Computes correspondences between 3D mesh and UV map
+NOTE: 3D mesh vertices can have multiple correspondences with UV vertices
+"""
+def mesh2UV(F, Ft):
+	m2uv = {v: set() for f in F for v in f}
+	for f, ft in zip(F, Ft):
+		for v, vt in zip(f, ft):
+			m2uv[v].add(vt)
+	# m2uv = {k:list(v) for k,v in m2uv.items()}
+	return m2uv
+
+
+def filter_m2u(m2u, to_keep):
+    for key in m2u:
+        if len(m2u[key]) == 2:
+            v = list(m2u[key])
+            if v[0] in to_keep:
+                m2u[key].remove(v[1])
+            else:
+                m2u[key].remove(v[0])
+    return m2u
+
+
+def convert_to_numbers(m2u):
+    for key in m2u:
+        m2u[key] = m2u[key].pop()
+    return m2u
+
+
+def get_masked_index(mask, index):
+    m = np.zeros((mask.shape[0]), dtype=bool)
+    m[index] = True
+    m = m[mask]
+    new_index = np.argmax(m)
+    return new_index
+
+
+def get_closest_point(verts_1, verts_2):
+    dist = sp.spatial.distance.cdist(verts_1, verts_2)
+    closest = np.argmin(dist, axis=-1)
+    return closest
+
+def get_normal_edge(p1, p2):
+    d = p2 - p1
+    return np.array([-d[1], d[0]])
+
+def get_normal_vertex(p1, p2, p3):
+    n1 = get_normal_edge(p1, p2)
+    n2 = get_normal_edge(p2, p3)
+    n = (n1 + n2)
+    n = n / np.linalg.norm(n) / 20
+    return n
+
+
+def get_distance(p1, p2, points):
+    d = np.array([np.linalg.norm(np.cross(p2 - p1, point - p1)) / np.linalg.norm(p2 - p1) for point in points])
+    return d
+
+
+def get_normals(verts):
+    v = verts.copy()
+    v = np.concatenate([v[-1][np.newaxis, :], v, v[0][np.newaxis, :]])
+    normals = np.zeros_like(verts)
+    for i in range(1, v.shape[0] - 1):
+        normals[i - 1] = get_normal_vertex(v[i - 1], v[i], v[i + 1])
+    return normals
+
+
+def get_closest_points_normals(verts1, verts2):
+    dist_2 = sp.spatial.distance.cdist(verts1, verts2)
+    normals = get_normals(verts1)
+
+    # fig = plt.figure(figsize=(15, 15))
+    # ax = fig.add_subplot(111)
+    # ax.set_xlim(0, 1)
+    # ax.set_ylim(0, 1)
+    # ax.set_aspect('equal')
+    # ax.scatter(verts1[:, 0], verts1[:, 1])
+    # for p1, p2 in zip(verts1, verts1 + normals):
+    #     ax.plot([p1[0], p2[0]], [p1[1], p2[1]])
+    # plt.show()
+
+    closest = []
+    i = 0
+    for p1, p2 in zip(verts1, verts1 + normals):
+        dist = get_distance(p1, p2, verts2)
+        closest.append(np.argmin(dist + dist_2[i, :]))
+        i += 1
+
+    return np.array(closest)
+
+
+def back():
+    # load data
+    data_body = np.load("Blender/uv_data_1.npz", allow_pickle=True)
+    data_back = np.load("Blender/uv_data_plane_1_2.npz", allow_pickle=True)
+    
+    # translate back uv to match body
+    uv_back = data_back.f.Vt
+    # uv_back[:, 0] = uv_back[:, 0] + 0.1
+    # uv_back[:, 1] = uv_back[:, 1] + 0.42
+
+    # find front and back indices in uv
+    uv_verts = np.append(data_body.f.Vt, np.zeros((data_body.f.Vt.shape[0], 1)), axis=-1)
+    m = tm.Trimesh(uv_verts, data_body.f.Ft, process=False, validate=False)
+    front, back = tm.graph.connected_components(m.edges, min_len=1, nodes=None, engine=None)
+
+    # get mesh to uv mapping for the back vertices
+    m2u_body = mesh2UV(data_body.f.F, data_body.f.Ft)
+    m2u_body = filter_m2u(m2u_body, set(back))
+    m2u_body = convert_to_numbers(m2u_body)
+
+    m2u_back = mesh2UV(data_back.f.F, data_back.f.Ft)
+    m2u_back = convert_to_numbers(m2u_back)
+
+    # create mask for filtering vertices belonging to the back
+    body_back_mask = np.zeros((data_body.f.Vt.shape[0]), dtype=bool)
+    body_back_mask[back] = True
+
+    # define which points should match after the transform on the mesh
+    matching_points = np.array([[3284, 3], [9557, 13], [9278, 15], [8095, 1], [10597, 10], [10602, 17], [7114, 9], [6989, 16], [6983, 11], [4219, 0], [5628, 14], [5908, 12], [1447, 2], [1086, 5], [21, 8], [33, 4], [1888, 7], [2931, 6]])
+    
+    # map matching point to uv coordinates
+    for i in range(matching_points.shape[0]):
+        matching_points[i, :] = [get_masked_index(body_back_mask, m2u_body[matching_points[i,0]]), m2u_back[matching_points[i, 1]]]
+
+    x = data_body.f.Vt[body_back_mask, :]
+    y = uv_back
+
+    fig = plt.figure(figsize=(15, 15))
+    ax = fig.add_subplot(121)
+    ax.set_xlim(0, 1)
+    ax.set_ylim(0, 1)
+    ax.set_aspect('equal')
+    ax.scatter(x[:, 0], x[:, 1])
+    if y is not None:
+        ax.scatter(y[:, 0], y[:, 1])
+    for pair in matching_points:
+        ax.scatter([x[pair[0], 0], y[pair[1], 0]], [x[pair[0], 1], y[pair[1], 1]], s=50)
+
+    # get transformation from body to back contour
+    source =  x[matching_points[:, 0], :]
+    destination = y[matching_points[:, 1], :]
+    trans = tps.TPS(source, destination)
+    transformed = trans(data_body.f.Vt[back, :])
+
+    ax = fig.add_subplot(122)
+    ax.set_xlim(0, 1)
+    ax.set_ylim(0, 1)
+    ax.set_aspect('equal')
+    ax.scatter(transformed[:, 0], transformed[:, 1])
+    if y is not None:
+        ax.scatter(y[:, 0], y[:, 1])
+    for pair in matching_points:
+        ax.scatter([transformed[pair[0], 0], y[pair[1], 0]], [transformed[pair[0], 1], y[pair[1], 1]], s=50)
+
+    plt.show()
+
+    data_back_high = np.load("Blender/uv_data_plane_1_high_res.npz", allow_pickle=True)
+    y_high = data_back_high.f.Vt
+    # y_high[:, 0] = y_high[:, 0] + 0.1
+    # y_high[:, 1] = y_high[:, 1] + 0.42
+    
+    closest_points = get_closest_points_normals(y_high, transformed)
+    # closest_points = get_closest_point(y_high, transformed)
+    unique_mask = np.zeros_like(closest_points, dtype=bool)
+    unique_mask[np.unique(closest_points, return_index=True)[1]] = True
+    matching_points = np.array([closest_points[unique_mask], np.array(range(y_high.shape[0]))[unique_mask]]).T
+    # matching_points_2 = []
+    # for i, p in enumerate(matching_points):
+    #     if i % 10 == 0:
+    #         matching_points_2.append(p)
+    # matching_points = np.array(matching_points_2)
+
+    source_2 = transformed[matching_points[:, 0], :]
+    destination_2 = y_high[matching_points[:, 1], :]
+    trans_2 = tps.TPS(source_2, destination_2)
+    transformed_2 = trans_2(transformed)
+
+
+
+    fig = plt.figure(figsize=(15, 15))
+    ax = fig.add_subplot(121)
+    ax.set_xlim(0, 1)
+    ax.set_ylim(0, 1)
+    ax.set_aspect('equal')
+    ax.scatter(source_2[:, 0], source_2[:, 1])
+    ax.scatter(destination_2[:, 0], destination_2[:, 1])
+    for i, pair in enumerate(zip(source_2, destination_2)):
+        if i % 5 == 0:
+            s, d = pair
+            ax.scatter([s[0], d[0]], [s[1], d[1]])
+
+    ax = fig.add_subplot(122)
+    ax.set_xlim(0, 1)
+    ax.set_ylim(0, 1)
+    ax.set_aspect('equal')
+    ax.scatter(transformed_2[:, 0], transformed_2[:, 1])
+    ax.scatter(y_high[:, 0], y_high[:, 1])
+    plt.show()
+
+    return transformed_2
+
+
+def match_back_to_front(front_uv):
+    # load data
+    data_body = np.load("Blender/uv_data_smpl_2.npz", allow_pickle=True)
+    # data_back = np.load("Blender/uv_data_plane_1_2.npz", allow_pickle=True)
+    
+    # translate back uv to match body
+    # uv_back = data_back.f.Vt
+    # uv_back[:, 0] = uv_back[:, 0] + 0.1
+    # uv_back[:, 1] = uv_back[:, 1] + 0.42
+
+    # find front and back indices in uv
+    uv_verts = np.append(data_body.f.Vt, np.zeros((data_body.f.Vt.shape[0], 1)), axis=-1)
+    m = tm.Trimesh(uv_verts, data_body.f.Ft, process=False, validate=False)
+    front, back = tm.graph.connected_components(m.edges, min_len=1, nodes=None, engine=None)
+
+    # get mesh to uv mapping for the back vertices
+    m2u_body = mesh2UV(data_body.f.F, data_body.f.Ft)
+    m2u_body = filter_m2u(m2u_body, set(back))
+    m2u_body = convert_to_numbers(m2u_body)
+
+    # m2u_back = mesh2UV(data_back.f.F, data_back.f.Ft)
+    # m2u_back = convert_to_numbers(m2u_back)
+
+    # create mask for filtering vertices belonging to the back
+    body_back_mask = np.zeros((data_body.f.Vt.shape[0]), dtype=bool)
+    body_back_mask[back] = True
+    body_front_mask = np.zeros((data_body.f.Vt.shape[0]), dtype=bool)
+    body_front_mask[front] = True
+
+    # define which points should match after the transform on the mesh
+    # matching_points = np.array([[3284, 3], [9557, 13], [9278, 15], [8095, 1], [10597, 10], [10602, 17], [7114, 9], [6989, 16], [6983, 11], [4219, 0], [5628, 14], [5908, 12], [1447, 2], [1086, 5], [21, 8], [33, 4], [1888, 7], [2931, 6]])
+    
+    # # map matching point to uv coordinates
+    # for i in range(matching_points.shape[0]):
+    #     matching_points[i, :] = [get_masked_index(body_back_mask, m2u_body[matching_points[i,0]]), m2u_back[matching_points[i, 1]]]
+
+    m2u_body_all = mesh2UV(data_body.f.F, data_body.f.Ft)
+    m2u_body_all = m2uv_array(m2u_body_all)
+    border_pairs = m2u_body_all[:, 0] != m2u_body_all[:, 1]
+    border_pairs = m2u_body_all[border_pairs, :]
+    matching_points = border_pairs
+    matching_points_sorted = []
+    for p in matching_points:
+        if p[0] in front:
+            matching_points_sorted.append([p[0], p[1]])
+        elif p[0] in back:
+            matching_points_sorted.append([p[1], p[0]])
+        else:
+            assert False
+    # add neck V shape control point
+    matching_points_sorted.append([m2u_body_all[10228][0], m2u_body_all[33][0]])
+    matching_points = np.array(matching_points_sorted)
+    mm = matching_points.copy()
+
+    for i in range(matching_points.shape[0]):
+        matching_points[i, :] = [get_masked_index(body_back_mask, matching_points[i,1]), get_masked_index(body_front_mask, matching_points[i,0])]
+
+    temp = np.zeros_like(data_body.f.Vt)
+    temp[front, :] = front_uv
+
+    x = data_body.f.Vt[body_back_mask, :]
+    y = temp[body_front_mask, :] # front_uv # 
+
+    fig = plt.figure(figsize=(15, 15))
+    ax = fig.add_subplot(121)
+    ax.set_xlim(0, 1)
+    ax.set_ylim(0, 1)
+    ax.set_aspect('equal')
+    ax.scatter(x[:, 0], x[:, 1])
+    if y is not None:
+        ax.scatter(y[:, 0], y[:, 1])
+    for pair in matching_points:
+        ax.scatter([x[pair[0], 0], y[pair[1], 0]], [x[pair[0], 1], y[pair[1], 1]], s=50)
+
+    # get transformation from body to back contour
+    source =  x[matching_points[:, 0], :]
+    destination = y[matching_points[:, 1], :]
+    trans = tps.TPS(source, destination)
+    transformed = trans(data_body.f.Vt[back, :])
+
+    ax = fig.add_subplot(122)
+    ax.set_xlim(0, 1)
+    ax.set_ylim(0, 1)
+    ax.set_aspect('equal')
+    ax.scatter(transformed[:, 0], transformed[:, 1])
+    if y is not None:
+        ax.scatter(y[:, 0], y[:, 1])
+    for pair in matching_points:
+        ax.scatter([transformed[pair[0], 0], y[pair[1], 0]], [transformed[pair[0], 1], y[pair[1], 1]], s=50)
+
+    plt.show()
+
+
+    return transformed
+
+
+def generate_mask():
+    import cv2
+    res = 1024
+    front = np.load("Blender/uv_data_plane_0_high_res_2.npz", allow_pickle=True)
+    back = np.load("Blender/uv_data_plane_1_high_res_2.npz", allow_pickle=True)
+    points = np.round(np.concatenate((front.f.Vt, back.f.Vt)) * res).astype(np.int32)
+    img = np.zeros((res, res, 3))
+    img[points[:, 0], points[:, 1], :] = 255
+    cv2.imwrite('MeshUtils/mask_1024.png', img)
+
+    # do paint now
+
+    img = cv2.imread('MeshUtils/mask_1024_2.png')
+    img = img[:, :, 0] == 255
+    np.save('MeshUtils/mask_1024.npy', img)
+
+def generate_mask_2():
+    import cv2
+    res = 256
+    front = np.load("Blender/uv_data_plane_stretched_front_high_extra.npz", allow_pickle=True)
+    back = np.load("Blender/uv_data_plane_stretched_front_high_extra.npz", allow_pickle=True)
+    points = np.floor(np.concatenate((front.f.Vt, back.f.Vt + [1, 0])) * res).astype(np.int32)
+    img = np.zeros((res * 2, res, 3))
+    img[points[:, 0], points[:, 1], :] = 255
+    cv2.imwrite('MeshUtils/mask_256x512.png', img)
+
+    # do paint now
+
+    img = cv2.imread('MeshUtils/mask_256x512.png')
+    img = img[:, :, 0] > 240
+    np.save('MeshUtils/mask_256x512.npy', img)
+
+    pass
+
+
+def generate_uv_mask_border():
+    import cv2
+    mask = np.load('MeshUtils/mask_256x512.npy')
+    uv_res = 256
+    mask_border = (cv2.resize(mask * 255.0, (uv_res, uv_res * 2)) / 255).astype(bool)[:, :, np.newaxis]
+    mask_border = cv2.erode(mask_border.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=1).astype(bool)
+
+    mask[mask_border == True] = 0
+    # plt.imshow(mask)
+    idx_front = np.transpose(np.vstack(np.nonzero(mask[:256, :])), [1,0])
+
+    test = np.zeros_like(mask)
+    test[idx_front[:,0], idx_front[:,1]] = 1
+    idx_back = idx_front.copy()
+    idx_back[:, 0] += 256
+    test[idx_back[:,0], idx_back[:,1]] = 1
+
+
+    pass
+
+
+
+def front():
+    # load data
+    data_body = np.load("custom_data/uv_map/uv_data_smpl_2.npz", allow_pickle=True)
+    data_front = np.load("custom_data/uv_map/uv_data_plane_stretched_front_low.npz", allow_pickle=True)
+    
+    # translate back uv to match body
+    uv_front = data_front.f.Vt
+    # uv_back[:, 0] = uv_back[:, 0] + 0.1
+    # uv_back[:, 1] = uv_back[:, 1] + 0.42
+
+    # find front and back indices in uv
+    uv_verts = np.append(data_body.f.Vt, np.zeros((data_body.f.Vt.shape[0], 1)), axis=-1)
+    m = tm.Trimesh(uv_verts, data_body.f.Ft, process=False, validate=False)
+    front, back = tm.graph.connected_components(m.edges, min_len=1, nodes=None, engine=None)
+
+    # get mesh to uv mapping for the back vertices
+    m2u_body = mesh2UV(data_body.f.F, data_body.f.Ft)
+    m2u_body = filter_m2u(m2u_body, set(front))
+    m2u_body = convert_to_numbers(m2u_body)
+
+    m2u_front = mesh2UV(data_front.f.F, data_front.f.Ft)
+    m2u_front = convert_to_numbers(m2u_front)
+
+    # create mask for filtering vertices belonging to the back
+    body_front_mask = np.zeros((data_body.f.Vt.shape[0]), dtype=bool)
+    body_front_mask[front] = True
+
+    # define which points should match after the transform on the mesh
+    matching_points = np.array([ [9557, 12], [3284, 2], [3053, 5], [1888, 8], [10228, 4], [21, 7], [1209, 6], [1447, 3], [5908, 13], [4617, 15], [4219, 1], [6983, 10], [6989, 17], [7114, 9], [10602, 16], [10597, 11], [7884, 0], [8280, 14] ])
+    # matching_points = np.array([ [9557, 12], [3284, 2], [3053, 5], [1888, 8], [10228, 4], [21, 7], [1209, 6], [1447, 3], [5908, 13], [4793, 15], [4219, 1], [6983, 10], [6989, 17], [7114, 9], [10602, 16], [10597, 11], [7884, 0], [8450, 14], [13566, 18], [11767, 19], [8668, 20], [5012, 21], [12716, 22], [10917, 23], [8280, 24], [4617, 25], [2760, 26], [911, 27] ])
+
+    # map matching point to uv coordinates
+    for i in range(matching_points.shape[0]):
+        matching_points[i, :] = [get_masked_index(body_front_mask, m2u_body[matching_points[i,0]]), m2u_front[matching_points[i, 1]]]
+
+    x = data_body.f.Vt[body_front_mask, :]
+    y = uv_front
+
+    fig = plt.figure(figsize=(15, 15))
+    ax = fig.add_subplot(121)
+    # ax.set_xlim(0, 1)
+    # ax.set_ylim(0, 1)
+    ax.set_aspect('equal')
+    ax.scatter(x[:, 0], x[:, 1])
+    if y is not None:
+        ax.scatter(y[:, 0], y[:, 1])
+    for pair in matching_points:
+        ax.scatter([x[pair[0], 0], y[pair[1], 0]], [x[pair[0], 1], y[pair[1], 1]], s=50)
+
+    # get transformation from body to front contour
+    source = x[matching_points[:, 0], :]
+    destination = y[matching_points[:, 1], :]
+    trans = tps.TPS(source, destination)
+    transformed = trans(data_body.f.Vt[front, :])
+
+    # transformed = (transformed - transformed.min()) / (transformed.max() - transformed.min())
+
+    ax = fig.add_subplot(122)
+    # ax.set_xlim(0, 1)
+    # ax.set_ylim(0, 1)
+    ax.set_aspect('equal')
+    ax.scatter(transformed[:, 0], transformed[:, 1])
+    if y is not None:
+        ax.scatter(y[:, 0], y[:, 1])
+    for pair in matching_points:
+        ax.scatter([transformed[pair[0], 0], y[pair[1], 0]], [transformed[pair[0], 1], y[pair[1], 1]], s=50)
+
+    plt.show()
+
+    data_front_high = np.load("custom_data/uv_map/uv_data_plane_stretched_front_high.npz", allow_pickle=True)
+    y_high = data_front_high.f.Vt
+    # y_high[:, 0] = y_high[:, 0] + 0.1
+    # y_high[:, 1] = y_high[:, 1] + 0.42
+    
+    closest_points = get_closest_points_normals(y_high, transformed)
+    # closest_points = get_closest_point(y_high, transformed)
+    unique_mask = np.zeros_like(closest_points, dtype=bool)
+    unique_mask[np.unique(closest_points, return_index=True)[1]] = True
+    matching_points = np.array([closest_points[unique_mask], np.array(range(y_high.shape[0]))[unique_mask]]).T
+    # matching_points_2 = []
+    # for i, p in enumerate(matching_points):
+    #     if i % 10 == 0:
+    #         matching_points_2.append(p)
+    # matching_points = np.array(matching_points_2)
+
+    source_2 = transformed[matching_points[:, 0], :]
+    destination_2 = y_high[matching_points[:, 1], :]
+    trans_2 = tps.TPS(source_2, destination_2)
+    transformed_2 = trans_2(transformed)
+
+
+
+    fig = plt.figure(figsize=(15, 15))
+    ax = fig.add_subplot(121)
+    ax.set_xlim(0, 1)
+    ax.set_ylim(0, 1)
+    ax.set_aspect('equal')
+    ax.scatter(source_2[:, 0], source_2[:, 1])
+    ax.scatter(destination_2[:, 0], destination_2[:, 1])
+    for i, pair in enumerate(zip(source_2, destination_2)):
+        if i % 5 == 0:
+            s, d = pair
+            ax.scatter([s[0], d[0]], [s[1], d[1]])
+
+    ax = fig.add_subplot(122)
+    ax.set_xlim(0, 1)
+    ax.set_ylim(0, 1)
+    ax.set_aspect('equal')
+    ax.scatter(transformed_2[:, 0], transformed_2[:, 1])
+    ax.scatter(y_high[:, 0], y_high[:, 1])
+    plt.show()
+
+    return transformed_2, mesh2UV(data_body.f.F, data_body.f.Ft), front, back
+
+
+def smoothing():
+    data_body = np.load("Blender/uv_data_smpl_2.npz", allow_pickle=True)
+    uv_map = np.load('MeshUtils\stretched_uv_map_3.npy')
+    verts = np.concatenate((uv_map, np.ones((15068, 1))), axis=-1)
+    m = tm.Trimesh(verts, data_body.f.Ft, process=False, validate=False)
+
+
+    m2u_body_all = mesh2UV(data_body.f.F, data_body.f.Ft)
+    m2u_body_all = m2uv_array(m2u_body_all)
+    border_pairs = m2u_body_all[:, 0] != m2u_body_all[:, 1]
+    border_pairs = m2u_body_all[border_pairs, :]
+    pinned_vertices = border_pairs.reshape(-1)
+
+    lap = tm.smoothing.laplacian_calculation(m, equal_weight=False, pinned_vertices=pinned_vertices)
+    # tm.smoothing.filter_laplacian(m, lamb=0.5, iterations=100, volume_constraint=False, laplacian_operator=lap)
+    # tm.smoothing.filter_mut_dif_laplacian(m, lamb=0.5, iterations=100, volume_constraint=False, laplacian_operator=lap)
+    # tm.smoothing.filter_humphrey(m, iterations=100, laplacian_operator=lap)
+    tm.smoothing.filter_taubin(m, iterations=100, laplacian_operator=lap)
+
+
+    x = uv_map
+    y = m.vertices[:, :2]
+
+    fig = plt.figure(figsize=(15, 15))
+    ax = fig.add_subplot(211)
+    ax.set_xlim(0, 2)
+    ax.set_ylim(0, 1)
+    # ax.set_aspect('equal')
+    ax.scatter(x[:, 0], x[:, 1])
+
+    ax = fig.add_subplot(212)
+    ax.set_xlim(0, 2)
+    ax.set_ylim(0, 1)
+    # ax.set_aspect('equal')
+    ax.scatter(y[:, 0], y[:, 1])
+    plt.show()
+
+
+    pass
+
+
+
+if __name__ == '__main__':
+    # generate_mask_2()
+    # smoothing()
+    # generate_uv_mask_border()
+
+    front_uv, m2uv, front_mask, back_mask = front()
+    back_uv = match_back_to_front(front_uv)
+    # back_uv = back()
+
+    final = np.zeros((15068, 2))
+    final[front_mask] = front_uv
+    final[back_mask] = back_uv + [1, 0]
+    np.save('MeshUtils\stretched_uv_map_3.npy', final)
+
+    pass

custom_utils/create_uv_map_2.py

@@ -0,0 +1,846 @@
+import os
+import sys
+import shutil
+import tqdm
+import random
+import math
+import collections
+from tps_custom import TPS
+import numpy as np
+import scipy as sp
+import trimesh as tm
+import cmapy as cm
+import scipy.io as sio
+import matplotlib.pyplot as plt
+import plotly.express as ex
+import plotly.graph_objects as go
+from scipy.spatial.distance import cdist
+
+
+
+def m2uv_array(m2uv, is_smpl_orig=False):
+    result = np.zeros((14475 if not is_smpl_orig else 3638, 2), dtype=np.int32)
+    for k, v in m2uv.items():
+        if len(v) == 2:
+            result[k] = list(v)
+        elif len(v) == 1:
+            x = v.pop()
+            result[k] = [x, x]
+        else:
+            assert False
+
+    return result
+
+
+def get_body():
+    smpl = SMPL('FeatureLines/superSMPL')
+    pose = np.zeros((1, 72), dtype=np.float32)
+    beta = np.zeros((1, 11), dtype=np.float32)
+    R = smpl.run_rodrigues(pose)
+    T, _, _ = smpl.run(R, beta, True)
+    return np.squeeze(T)
+
+
+def plot_3d(x, y=None):
+    fig = plt.figure(figsize=(15, 15))
+    ax = fig.add_subplot(111, projection='3d')
+    ax.set_xlim3d(-1, 1)
+    ax.set_ylim3d(-1, 1)
+    ax.set_zlim3d(-1, 1)
+    ax.set_xlabel("X")
+    ax.set_ylabel("Y")
+    ax.set_zlabel("Z")
+    ax.scatter(x[:, 0], x[:, 1], x[:, 2], s=5)
+    if y is not None:
+        ax.scatter(y[:, 0], y[:, 1], y[:, 2], s=5)
+    plt.show()
+
+def plot_uv(x, y=None):
+    fig = plt.figure(figsize=(15, 15))
+    ax = fig.add_subplot(111)
+    ax.set_xlim(0, 1)
+    ax.set_ylim(0, 1)
+    ax.scatter(x[:, 0], x[:, 1])
+    if y is not None:
+        ax.scatter(y[:, 0], y[:, 1])
+    plt.show()
+
+""" 
+Mesh to UV map
+Computes correspondences between 3D mesh and UV map
+NOTE: 3D mesh vertices can have multiple correspondences with UV vertices
+"""
+def mesh2UV(F, Ft):
+	m2uv = {v: set() for f in F for v in f}
+	for f, ft in zip(F, Ft):
+		for v, vt in zip(f, ft):
+			m2uv[v].add(vt)
+	# m2uv = {k:list(v) for k,v in m2uv.items()}
+	return m2uv
+
+
+def filter_m2u(m2u, to_keep):
+    for key in m2u:
+        if len(m2u[key]) == 2:
+            v = list(m2u[key])
+            if v[0] in to_keep:
+                m2u[key].remove(v[1])
+            else:
+                m2u[key].remove(v[0])
+    return m2u
+
+
+def convert_to_numbers(m2u):
+    for key in m2u:
+        m2u[key] = m2u[key].pop()
+    return m2u
+
+
+def get_masked_index(mask, index):
+    m = np.zeros((mask.shape[0]), dtype=bool)
+    m[index] = True
+    m = m[mask]
+    new_index = np.argmax(m)
+    return new_index
+
+
+def get_closest_point(verts_1, verts_2):
+    dist = sp.spatial.distance.cdist(verts_1, verts_2)
+    closest = np.argmin(dist, axis=-1)
+    return closest
+
+def get_normal_edge(p1, p2):
+    d = p2 - p1
+    return np.array([-d[1], d[0]])
+
+def get_normal_vertex(p1, p2, p3):
+    n1 = get_normal_edge(p1, p2)
+    n2 = get_normal_edge(p2, p3)
+    n = (n1 + n2)
+    n = n / np.linalg.norm(n) / 20
+    return n
+
+
+def get_distance(p1, p2, points):
+    d = np.array([np.linalg.norm(np.cross(p2 - p1, point - p1)) / np.linalg.norm(p2 - p1) for point in points])
+    return d
+
+
+def get_normals(verts):
+    v = verts.copy()
+    v = np.concatenate([v[-1][np.newaxis, :], v, v[0][np.newaxis, :]])
+    normals = np.zeros_like(verts)
+    for i in range(1, v.shape[0] - 1):
+        normals[i - 1] = get_normal_vertex(v[i - 1], v[i], v[i + 1])
+    return normals
+
+
+def get_closest_points_normals(verts1, verts2):
+    dist_2 = sp.spatial.distance.cdist(verts1, verts2)
+    normals = get_normals(verts1)
+
+    # fig = plt.figure(figsize=(15, 15))
+    # ax = fig.add_subplot(111)
+    # ax.set_xlim(0, 1)
+    # ax.set_ylim(0, 1)
+    # ax.set_aspect('equal')
+    # ax.scatter(verts1[:, 0], verts1[:, 1])
+    # for p1, p2 in zip(verts1, verts1 + normals):
+    #     ax.plot([p1[0], p2[0]], [p1[1], p2[1]])
+    # plt.show()
+
+    closest = []
+    i = 0
+    for p1, p2 in zip(verts1, verts1 + normals):
+        dist = get_distance(p1, p2, verts2)
+        closest.append(np.argmin(dist + dist_2[i, :]))
+        i += 1
+
+    return np.array(closest)
+
+
+def back():
+    # load data
+    data_body = np.load("Blender/uv_data_1.npz", allow_pickle=True)
+    data_back = np.load("Blender/uv_data_plane_1_2.npz", allow_pickle=True)
+    
+    # translate back uv to match body
+    uv_back = data_back.f.Vt
+    # uv_back[:, 0] = uv_back[:, 0] + 0.1
+    # uv_back[:, 1] = uv_back[:, 1] + 0.42
+
+    # find front and back indices in uv
+    uv_verts = np.append(data_body.f.Vt, np.zeros((data_body.f.Vt.shape[0], 1)), axis=-1)
+    m = tm.Trimesh(uv_verts, data_body.f.Ft, process=False, validate=False)
+    front, back = tm.graph.connected_components(m.edges, min_len=1, nodes=None, engine=None)
+
+    # get mesh to uv mapping for the back vertices
+    m2u_body = mesh2UV(data_body.f.F, data_body.f.Ft)
+    m2u_body = filter_m2u(m2u_body, set(back))
+    m2u_body = convert_to_numbers(m2u_body)
+
+    m2u_back = mesh2UV(data_back.f.F, data_back.f.Ft)
+    m2u_back = convert_to_numbers(m2u_back)
+
+    # create mask for filtering vertices belonging to the back
+    body_back_mask = np.zeros((data_body.f.Vt.shape[0]), dtype=bool)
+    body_back_mask[back] = True
+
+    # define which points should match after the transform on the mesh
+    matching_points = np.array([[3284, 3], [9557, 13], [9278, 15], [8095, 1], [10597, 10], [10602, 17], [7114, 9], [6989, 16], [6983, 11], [4219, 0], [5628, 14], [5908, 12], [1447, 2], [1086, 5], [21, 8], [33, 4], [1888, 7], [2931, 6]])
+    
+    # map matching point to uv coordinates
+    for i in range(matching_points.shape[0]):
+        matching_points[i, :] = [get_masked_index(body_back_mask, m2u_body[matching_points[i,0]]), m2u_back[matching_points[i, 1]]]
+
+    x = data_body.f.Vt[body_back_mask, :]
+    y = uv_back
+
+    fig = plt.figure(figsize=(15, 15))
+    ax = fig.add_subplot(121)
+    ax.set_xlim(0, 1)
+    ax.set_ylim(0, 1)
+    ax.set_aspect('equal')
+    ax.scatter(x[:, 0], x[:, 1])
+    if y is not None:
+        ax.scatter(y[:, 0], y[:, 1])
+    for pair in matching_points:
+        ax.scatter([x[pair[0], 0], y[pair[1], 0]], [x[pair[0], 1], y[pair[1], 1]], s=50)
+
+    # get transformation from body to back contour
+    source =  x[matching_points[:, 0], :]
+    destination = y[matching_points[:, 1], :]
+    trans = tps.TPS(source, destination)
+    transformed = trans(data_body.f.Vt[back, :])
+
+    ax = fig.add_subplot(122)
+    ax.set_xlim(0, 1)
+    ax.set_ylim(0, 1)
+    ax.set_aspect('equal')
+    ax.scatter(transformed[:, 0], transformed[:, 1])
+    if y is not None:
+        ax.scatter(y[:, 0], y[:, 1])
+    for pair in matching_points:
+        ax.scatter([transformed[pair[0], 0], y[pair[1], 0]], [transformed[pair[0], 1], y[pair[1], 1]], s=50)
+
+    plt.show()
+
+    data_back_high = np.load("Blender/uv_data_plane_1_high_res.npz", allow_pickle=True)
+    y_high = data_back_high.f.Vt
+    # y_high[:, 0] = y_high[:, 0] + 0.1
+    # y_high[:, 1] = y_high[:, 1] + 0.42
+    
+    closest_points = get_closest_points_normals(y_high, transformed)
+    # closest_points = get_closest_point(y_high, transformed)
+    unique_mask = np.zeros_like(closest_points, dtype=bool)
+    unique_mask[np.unique(closest_points, return_index=True)[1]] = True
+    matching_points = np.array([closest_points[unique_mask], np.array(range(y_high.shape[0]))[unique_mask]]).T
+    # matching_points_2 = []
+    # for i, p in enumerate(matching_points):
+    #     if i % 10 == 0:
+    #         matching_points_2.append(p)
+    # matching_points = np.array(matching_points_2)
+
+    source_2 = transformed[matching_points[:, 0], :]
+    destination_2 = y_high[matching_points[:, 1], :]
+    trans_2 = tps.TPS(source_2, destination_2)
+    transformed_2 = trans_2(transformed)
+
+
+
+    fig = plt.figure(figsize=(15, 15))
+    ax = fig.add_subplot(121)
+    ax.set_xlim(0, 1)
+    ax.set_ylim(0, 1)
+    ax.set_aspect('equal')
+    ax.scatter(source_2[:, 0], source_2[:, 1])
+    ax.scatter(destination_2[:, 0], destination_2[:, 1])
+    for i, pair in enumerate(zip(source_2, destination_2)):
+        if i % 5 == 0:
+            s, d = pair
+            ax.scatter([s[0], d[0]], [s[1], d[1]])
+
+    ax = fig.add_subplot(122)
+    ax.set_xlim(0, 1)
+    ax.set_ylim(0, 1)
+    ax.set_aspect('equal')
+    ax.scatter(transformed_2[:, 0], transformed_2[:, 1])
+    ax.scatter(y_high[:, 0], y_high[:, 1])
+    plt.show()
+
+    return transformed_2
+
+
+def match_back_to_front(front_uv):
+    # load data
+    data_body = np.load("Blender/uv_data_smpl_2.npz", allow_pickle=True)
+    # data_back = np.load("Blender/uv_data_plane_1_2.npz", allow_pickle=True)
+    
+    # translate back uv to match body
+    # uv_back = data_back.f.Vt
+    # uv_back[:, 0] = uv_back[:, 0] + 0.1
+    # uv_back[:, 1] = uv_back[:, 1] + 0.42
+
+    # find front and back indices in uv
+    uv_verts = np.append(data_body.f.Vt, np.zeros((data_body.f.Vt.shape[0], 1)), axis=-1)
+    m = tm.Trimesh(uv_verts, data_body.f.Ft, process=False, validate=False)
+    front, back = tm.graph.connected_components(m.edges, min_len=1, nodes=None, engine=None)
+
+    # get mesh to uv mapping for the back vertices
+    m2u_body = mesh2UV(data_body.f.F, data_body.f.Ft)
+    m2u_body = filter_m2u(m2u_body, set(back))
+    m2u_body = convert_to_numbers(m2u_body)
+
+    # m2u_back = mesh2UV(data_back.f.F, data_back.f.Ft)
+    # m2u_back = convert_to_numbers(m2u_back)
+
+    # create mask for filtering vertices belonging to the back
+    body_back_mask = np.zeros((data_body.f.Vt.shape[0]), dtype=bool)
+    body_back_mask[back] = True
+    body_front_mask = np.zeros((data_body.f.Vt.shape[0]), dtype=bool)
+    body_front_mask[front] = True
+
+    # define which points should match after the transform on the mesh
+    # matching_points = np.array([[3284, 3], [9557, 13], [9278, 15], [8095, 1], [10597, 10], [10602, 17], [7114, 9], [6989, 16], [6983, 11], [4219, 0], [5628, 14], [5908, 12], [1447, 2], [1086, 5], [21, 8], [33, 4], [1888, 7], [2931, 6]])
+    
+    # # map matching point to uv coordinates
+    # for i in range(matching_points.shape[0]):
+    #     matching_points[i, :] = [get_masked_index(body_back_mask, m2u_body[matching_points[i,0]]), m2u_back[matching_points[i, 1]]]
+
+    m2u_body_all = mesh2UV(data_body.f.F, data_body.f.Ft)
+    m2u_body_all = m2uv_array(m2u_body_all)
+    border_pairs = m2u_body_all[:, 0] != m2u_body_all[:, 1]
+    border_pairs = m2u_body_all[border_pairs, :]
+    matching_points = border_pairs
+    matching_points_sorted = []
+    for p in matching_points:
+        if p[0] in front:
+            matching_points_sorted.append([p[0], p[1]])
+        elif p[0] in back:
+            matching_points_sorted.append([p[1], p[0]])
+        else:
+            assert False
+    # add neck V shape control point
+    matching_points_sorted.append([m2u_body_all[10228][0], m2u_body_all[33][0]])
+    matching_points = np.array(matching_points_sorted)
+    mm = matching_points.copy()
+
+    for i in range(matching_points.shape[0]):
+        matching_points[i, :] = [get_masked_index(body_back_mask, matching_points[i,1]), get_masked_index(body_front_mask, matching_points[i,0])]
+
+    temp = np.zeros_like(data_body.f.Vt)
+    temp[front, :] = front_uv
+
+    x = data_body.f.Vt[body_back_mask, :]
+    y = temp[body_front_mask, :] # front_uv # 
+
+    fig = plt.figure(figsize=(15, 15))
+    ax = fig.add_subplot(121)
+    ax.set_xlim(0, 1)
+    ax.set_ylim(0, 1)
+    ax.set_aspect('equal')
+    ax.scatter(x[:, 0], x[:, 1])
+    if y is not None:
+        ax.scatter(y[:, 0], y[:, 1])
+    for pair in matching_points:
+        ax.scatter([x[pair[0], 0], y[pair[1], 0]], [x[pair[0], 1], y[pair[1], 1]], s=50)
+
+    # get transformation from body to back contour
+    source =  x[matching_points[:, 0], :]
+    destination = y[matching_points[:, 1], :]
+    trans = tps.TPS(source, destination)
+    transformed = trans(data_body.f.Vt[back, :])
+
+    ax = fig.add_subplot(122)
+    ax.set_xlim(0, 1)
+    ax.set_ylim(0, 1)
+    ax.set_aspect('equal')
+    ax.scatter(transformed[:, 0], transformed[:, 1])
+    if y is not None:
+        ax.scatter(y[:, 0], y[:, 1])
+    for pair in matching_points:
+        ax.scatter([transformed[pair[0], 0], y[pair[1], 0]], [transformed[pair[0], 1], y[pair[1], 1]], s=50)
+
+    plt.show()
+
+
+    return transformed
+
+
+def generate_mask():
+    import cv2
+    res = 1024
+    front = np.load("Blender/uv_data_plane_0_high_res_2.npz", allow_pickle=True)
+    back = np.load("Blender/uv_data_plane_1_high_res_2.npz", allow_pickle=True)
+    points = np.round(np.concatenate((front.f.Vt, back.f.Vt)) * res).astype(np.int32)
+    img = np.zeros((res, res, 3))
+    img[points[:, 0], points[:, 1], :] = 255
+    cv2.imwrite('MeshUtils/mask_1024.png', img)
+
+    # do paint now
+
+    img = cv2.imread('MeshUtils/mask_1024_2.png')
+    img = img[:, :, 0] == 255
+    np.save('MeshUtils/mask_1024.npy', img)
+
+def generate_mask_2():
+    import cv2
+    res = 256
+    front = np.load("Blender/uv_data_plane_stretched_front_high_extra.npz", allow_pickle=True)
+    back = np.load("Blender/uv_data_plane_stretched_front_high_extra.npz", allow_pickle=True)
+    points = np.floor(np.concatenate((front.f.Vt, back.f.Vt + [1, 0])) * res).astype(np.int32)
+    img = np.zeros((res * 2, res, 3))
+    img[points[:, 0], points[:, 1], :] = 255
+    cv2.imwrite('MeshUtils/mask_256x512.png', img)
+
+    # do paint now
+
+    img = cv2.imread('MeshUtils/mask_256x512.png')
+    img = img[:, :, 0] > 240
+    np.save('MeshUtils/mask_256x512.npy', img)
+
+    pass
+
+
+def generate_uv_mask_border():
+    import cv2
+    mask = np.load('MeshUtils/mask_256x512.npy')
+    uv_res = 256
+    mask_border = (cv2.resize(mask * 255.0, (uv_res, uv_res * 2)) / 255).astype(bool)[:, :, np.newaxis]
+    mask_border = cv2.erode(mask_border.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=1).astype(bool)
+
+    mask[mask_border == True] = 0
+    # plt.imshow(mask)
+    idx_front = np.transpose(np.vstack(np.nonzero(mask[:256, :])), [1,0])
+
+    test = np.zeros_like(mask)
+    test[idx_front[:,0], idx_front[:,1]] = 1
+    idx_back = idx_front.copy()
+    idx_back[:, 0] += 256
+    test[idx_back[:,0], idx_back[:,1]] = 1
+
+
+    pass
+
+
+
+def front():
+    # load data
+    data_body = np.load("custom_data/uv_map/uv_data_smpl_2.npz", allow_pickle=True)
+    data_front = np.load("custom_data/uv_map/uv_data_plane_stretched_front_low.npz", allow_pickle=True)
+    
+    # translate back uv to match body
+    uv_front = data_front.f.Vt
+    # uv_back[:, 0] = uv_back[:, 0] + 0.1
+    # uv_back[:, 1] = uv_back[:, 1] + 0.42
+
+    # find front and back indices in uv
+    uv_verts = np.append(data_body.f.Vt, np.zeros((data_body.f.Vt.shape[0], 1)), axis=-1)
+    m = tm.Trimesh(uv_verts, data_body.f.Ft, process=False, validate=False)
+    front, back = tm.graph.connected_components(m.edges, min_len=1, nodes=None, engine=None)
+
+    # get mesh to uv mapping for the back vertices
+    m2u_body = mesh2UV(data_body.f.F, data_body.f.Ft)
+    m2u_body = filter_m2u(m2u_body, set(front))
+    m2u_body = convert_to_numbers(m2u_body)
+
+    m2u_front = mesh2UV(data_front.f.F, data_front.f.Ft)
+    m2u_front = convert_to_numbers(m2u_front)
+
+    # create mask for filtering vertices belonging to the back
+    body_front_mask = np.zeros((data_body.f.Vt.shape[0]), dtype=bool)
+    body_front_mask[front] = True
+
+    # define which points should match after the transform on the mesh
+    matching_points = np.array([ [9557, 12], [3284, 2], [3053, 5], [1888, 8], [10228, 4], [21, 7], [1209, 6], [1447, 3], [5908, 13], [4617, 15], [4219, 1], [6983, 10], [6989, 17], [7114, 9], [10602, 16], [10597, 11], [7884, 0], [8280, 14] ])
+    # matching_points = np.array([ [9557, 12], [3284, 2], [3053, 5], [1888, 8], [10228, 4], [21, 7], [1209, 6], [1447, 3], [5908, 13], [4793, 15], [4219, 1], [6983, 10], [6989, 17], [7114, 9], [10602, 16], [10597, 11], [7884, 0], [8450, 14], [13566, 18], [11767, 19], [8668, 20], [5012, 21], [12716, 22], [10917, 23], [8280, 24], [4617, 25], [2760, 26], [911, 27] ])
+
+    # map matching point to uv coordinates
+    for i in range(matching_points.shape[0]):
+        matching_points[i, :] = [get_masked_index(body_front_mask, m2u_body[matching_points[i,0]]), m2u_front[matching_points[i, 1]]]
+
+    x = data_body.f.Vt[body_front_mask, :]
+    y = uv_front
+
+    fig = plt.figure(figsize=(15, 15))
+    ax = fig.add_subplot(121)
+    # ax.set_xlim(0, 1)
+    # ax.set_ylim(0, 1)
+    ax.set_aspect('equal')
+    ax.scatter(x[:, 0], x[:, 1])
+    if y is not None:
+        ax.scatter(y[:, 0], y[:, 1])
+    for pair in matching_points:
+        ax.scatter([x[pair[0], 0], y[pair[1], 0]], [x[pair[0], 1], y[pair[1], 1]], s=50)
+
+    # get transformation from body to front contour
+    source = x[matching_points[:, 0], :]
+    destination = y[matching_points[:, 1], :]
+    trans = tps.TPS(source, destination)
+    transformed = trans(data_body.f.Vt[front, :])
+
+    # transformed = (transformed - transformed.min()) / (transformed.max() - transformed.min())
+
+    ax = fig.add_subplot(122)
+    # ax.set_xlim(0, 1)
+    # ax.set_ylim(0, 1)
+    ax.set_aspect('equal')
+    ax.scatter(transformed[:, 0], transformed[:, 1])
+    if y is not None:
+        ax.scatter(y[:, 0], y[:, 1])
+    for pair in matching_points:
+        ax.scatter([transformed[pair[0], 0], y[pair[1], 0]], [transformed[pair[0], 1], y[pair[1], 1]], s=50)
+
+    plt.show()
+
+    data_front_high = np.load("custom_data/uv_map/uv_data_plane_stretched_front_high.npz", allow_pickle=True)
+    y_high = data_front_high.f.Vt
+    # y_high[:, 0] = y_high[:, 0] + 0.1
+    # y_high[:, 1] = y_high[:, 1] + 0.42
+    
+    closest_points = get_closest_points_normals(y_high, transformed)
+    # closest_points = get_closest_point(y_high, transformed)
+    unique_mask = np.zeros_like(closest_points, dtype=bool)
+    unique_mask[np.unique(closest_points, return_index=True)[1]] = True
+    matching_points = np.array([closest_points[unique_mask], np.array(range(y_high.shape[0]))[unique_mask]]).T
+    # matching_points_2 = []
+    # for i, p in enumerate(matching_points):
+    #     if i % 10 == 0:
+    #         matching_points_2.append(p)
+    # matching_points = np.array(matching_points_2)
+
+    source_2 = transformed[matching_points[:, 0], :]
+    destination_2 = y_high[matching_points[:, 1], :]
+    trans_2 = tps.TPS(source_2, destination_2)
+    transformed_2 = trans_2(transformed)
+
+
+
+    fig = plt.figure(figsize=(15, 15))
+    ax = fig.add_subplot(121)
+    ax.set_xlim(0, 1)
+    ax.set_ylim(0, 1)
+    ax.set_aspect('equal')
+    ax.scatter(source_2[:, 0], source_2[:, 1])
+    ax.scatter(destination_2[:, 0], destination_2[:, 1])
+    for i, pair in enumerate(zip(source_2, destination_2)):
+        if i % 5 == 0:
+            s, d = pair
+            ax.scatter([s[0], d[0]], [s[1], d[1]])
+
+    ax = fig.add_subplot(122)
+    ax.set_xlim(0, 1)
+    ax.set_ylim(0, 1)
+    ax.set_aspect('equal')
+    ax.scatter(transformed_2[:, 0], transformed_2[:, 1])
+    ax.scatter(y_high[:, 0], y_high[:, 1])
+    plt.show()
+
+    return transformed_2, mesh2UV(data_body.f.F, data_body.f.Ft), front, back
+
+
+
+def faces2graph(F, v=None):
+    if v: v = set(v) # Only this verts
+    G = {}
+    for f in F:
+        for i in f:
+            if v and i not in v: continue
+            if not i in G:
+                G[i] = {}
+            for j in f:
+                if v and j not in v: continue
+                if i == j: continue
+                G[i][j] = 1
+    return G
+
+
+def custom_smoothing(mesh, edges, pinned, F, iterations=10):
+    # custom_scatter(mesh[pinned, :]).show()
+
+    mesh_orig = mesh.copy()
+    G = faces2graph(F)
+    for _ in tqdm.tqdm(range(iterations)):
+        edge_lengths = np.sqrt(((mesh[edges[:,0], :] - mesh[edges[:,1], :]) ** 2 ).sum(-1))
+
+        edge_lengths_m = np.zeros((15068, 15068), np.float32)
+        edge_lengths_m[edges[:,0], edges[:,1]] = edge_lengths
+        edge_lengths_m[edges[:,1], edges[:,0]] = edge_lengths
+
+        # mean_edge = edge_lengths.mean()
+
+        # dist = cdist(mesh, mesh)
+        # for i in range(dist.shape[0]):
+        #     dist[i,i] = 10000
+
+        # dist = sp.spatial.KDTree(mesh)
+
+        # closest_neighbour = edge_lengths_m.argmin(axis=0)
+        #     if edge_lengths_m[i, closest_neighbour[i]] > mean_edge:
+        #         continue
+
+        mesh_temp = mesh.copy()
+        for i in range(15068):
+            if i in pinned:
+                continue
+
+            # closest_neighbour = edge_lengths_m[i, :].argmin()
+            # if edge_lengths_m[i, closest_neighbour] > mean_edge:
+            #     continue
+
+            neighbours = list(G[i].keys())
+            neighbours_dist = edge_lengths_m[i, neighbours]
+            weights = neighbours_dist / neighbours_dist.sum()
+            weights = np.ones_like(weights) / weights.shape[0]
+            mesh_temp[i, :] = np.sum(mesh[neighbours, :] * weights[:, np.newaxis], axis=0)
+
+            # neighbours_dist, neighbours = dist.query(mesh[i,:], k=10, workers=10)
+            # # if neighbours[0] == 15068:
+            # #     continue
+            # neighbours_dist, neighbours = neighbours_dist[1:], neighbours[1:]
+            # weights = neighbours_dist / neighbours_dist.sum()            
+            # mesh_temp[i, :] = np.sum(mesh[neighbours, :] * weights[:, np.newaxis], axis=0)
+
+            # mesh_temp[i, :] = np.mean(mesh[neighbours, :], axis=0)
+
+            # print()
+
+            # closest_neighbour = edge_lengths_m[i, :].argmin()
+            # if edge_lengths_m[i, closest_neighbour] < mean_edge:
+            #     vector = mesh[i,:] - mesh[closest_neighbour,:]
+            #     vector = vector / np.sqrt(np.sum(vector[0] ** 2 + vector[1] ** 2 ))
+            #     to_move = vector / 1000
+            #     # to_move = vector / edge_lengths_m[i, closest_neighbour] * 100
+            #     mesh[i, :] -= to_move
+
+            mesh = mesh_temp
+
+
+    custom_scatter(mesh_orig, mesh).show()
+    print()
+
+    return mesh
+
+def custom_scatter(x, y=None):
+    data = [go.Scatter(
+                x=x[:,0],
+                y=x[:,1],
+                mode='markers',
+                marker=dict(
+                    size=3,
+                    line=dict(width=0)
+                ),
+            ),]
+
+    if y is not None:
+        data.append(
+            go.Scatter(
+                x=y[:,0],
+                y=y[:,1],
+                mode='markers',
+                marker=dict(
+                    size=3,
+                    line=dict(width=0)
+                ),
+            ),)
+
+    fig = go.Figure(data=data)
+
+    # 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],),)
+    #                     )
+
+    return fig
+
+
+def smoothing():
+    data_body = np.load("custom_data/uv_map/uv_data_smpl_2.npz", allow_pickle=True)
+    uv_map = np.load('custom_data/uv_map/stretched_uv_map_3.npy')
+    F = np.load("custom_utils/smpl_faces.npy").astype(np.int32)
+    verts = np.concatenate((uv_map, np.ones((15068, 1))), axis=-1)
+    m = tm.Trimesh(verts, data_body.f.Ft, process=False, validate=False)
+
+
+    m2u_body_all = mesh2UV(data_body.f.F, data_body.f.Ft)
+    m2u_body_all = m2uv_array(m2u_body_all)
+    border_pairs = m2u_body_all[:, 0] != m2u_body_all[:, 1]
+    border_pairs = m2u_body_all[border_pairs, :]
+    pinned_vertices = border_pairs.reshape(-1)
+    right_legs_bottom = np.nonzero(((verts[:,0] > 1).astype(np.int8) * (verts[:,0] < 1.03).astype(np.int8)).astype(bool))[0]
+    left_legs_bottom = np.nonzero(verts[:,0] < 0.025)[0]
+    top_arms = np.nonzero(verts[:,1] > 0.995)[0]
+    bottom_arms = np.nonzero(verts[:,1] < 0.005)[0]
+    pinned_vertices = np.concatenate((pinned_vertices, right_legs_bottom, left_legs_bottom, top_arms, bottom_arms))
+
+    x = custom_smoothing(np.asarray(m.vertices)[:,:2], np.asarray(m.edges), pinned_vertices, data_body.f.Ft)
+    # m = tm.Trimesh(x, data_body.f.Ft, process=False, validate=False)
+
+    lap = tm.smoothing.laplacian_calculation(m, equal_weight=False, pinned_vertices=pinned_vertices)
+    # tm.smoothing.filter_laplacian(m, lamb=0.5, implicit_time_integration=False, iterations=1000, volume_constraint=False, laplacian_operator=lap)
+    # tm.smoothing.filter_mut_dif_laplacian(m, lamb=0.9, iterations=100, volume_constraint=True, laplacian_operator=lap)
+    tm.smoothing.filter_humphrey(m, iterations=1000, laplacian_operator=lap, alpha=0.01, beta=1.0)
+    # tm.smoothing.filter_taubin(m, iterations=25, lamb=0.01, nu=0.0015, laplacian_operator=lap)
+
+
+    x = uv_map
+    y = m.vertices[:, :2]
+
+
+
+    fig = go.Figure(data=[
+            # go.Scatter(
+            #     x=x[:,0],
+            #     y=x[:,1],
+            #     mode='markers',
+            #     marker=dict(
+            #         size=3,
+            #         line=dict(width=0)
+            #     ),
+            # ),
+            go.Scatter(
+                x=y[:,0],
+                y=y[:,1],
+                mode='markers',
+                marker=dict(
+                    size=3,
+                    line=dict(width=0)
+                ),
+            ),
+        ])
+
+    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()
+
+    dist = cdist(y, y)
+    for i in range(dist.shape[0]):
+        dist[i,i] = 10000
+
+    res  = 1 / dist.min()
+
+    print()
+
+
+    pass
+
+
+def front_3dto2d():
+    # load data
+    data_body = np.load("custom_data/uv_map/uv_data_smpl_2.npz", allow_pickle=True)
+    data_front = np.load("custom_data/uv_map/uv_data_plane_stretched_front_low.npz", allow_pickle=True)
+    
+    laplacian_loss = sio.loadmat("custom_utils/ssmpl0.mat")['L'].tocoo()
+    edges = np.stack(np.nonzero(laplacian_loss.todense())).transpose()
+
+    # translate back uv to match body
+    uv_front = data_front.f.V
+    # uv_back[:, 0] = uv_back[:, 0] + 0.1
+    # uv_back[:, 1] = uv_back[:, 1] + 0.42
+
+    # find front and back indices in uv
+    uv_verts = np.append(data_body.f.Vt, np.zeros((data_body.f.Vt.shape[0], 1)), axis=-1)
+    m = tm.Trimesh(uv_verts, data_body.f.Ft, process=False, validate=False)
+    front, back = tm.graph.connected_components(m.edges, min_len=1, nodes=None, engine=None)
+
+    # get mesh to uv mapping for the back vertices
+    m2u_body = mesh2UV(data_body.f.F, data_body.f.Ft)
+    m2u_body = filter_m2u(m2u_body, set(front))
+    m2u_body = convert_to_numbers(m2u_body)
+
+    m2u_front = mesh2UV(data_front.f.F, data_front.f.Ft)
+    m2u_front = convert_to_numbers(m2u_front)
+
+    # create mask for filtering vertices belonging to the back
+    body_front_mask = np.zeros((data_body.f.Vt.shape[0]), dtype=bool)
+    body_front_mask[front] = True
+
+    # define which points should match after the transform on the mesh
+    matching_points = np.array([ [9557, 12], [3284, 2], [3053, 5], [1888, 8], [10228, 4], [21, 7], [1209, 6], [1447, 3], [5908, 13], [4617, 15], [4219, 1], [6983, 10], [6989, 17], [7114, 9], [10602, 16], [10597, 11], [7884, 0], [8280, 14] ])
+
+    # for i in range(matching_points.shape[0]):
+    #     matching_points[i, :] = [get_masked_index(body_front_mask, matching_points[i,0]),matching_points[i, 1]]
+
+
+    x = data_body.f.V
+    # x = data_body.f.V[body_front_mask, :]
+    y = uv_front
+
+    # get transformation from body to front contour
+    source = x[matching_points[:, 0], :]
+    destination = y[matching_points[:, 1], :]
+    trans = TPS(source, destination)
+    transformed = trans(data_body.f.V)
+
+    z = transformed
+    fig = go.Figure(data=[
+            # go.Scatter3d(
+            #     x=x[:,0],
+            #     y=x[:,1],
+            #     z=x[:,2],
+            # ),
+            go.Scatter3d(
+                x=y[:,0],
+                y=y[:,1],
+                z=y[:,2],
+                mode='markers',
+                marker=dict(
+                    size=10,
+                    line=dict(width=0)
+                ),
+            ),
+            go.Scatter3d(
+                x=z[:,0],
+                y=z[:,1],
+                z=z[:,2],
+                mode='markers',
+                marker=dict(
+                    size=10,
+                    line=dict(width=0)
+                ),
+
+            ),
+        ])
+
+    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()
+
+
+if __name__ == '__main__':
+    # generate_mask_2()
+    smoothing()
+    # generate_uv_mask_border()
+
+    # front_uv, m2uv, front_mask, back_mask = front_3dto2d()
+
+
+    # back_uv = match_back_to_front(front_uv)
+    # # back_uv = back()
+
+    # final = np.zeros((15068, 2))
+    # final[front_mask] = front_uv
+    # final[back_mask] = back_uv + [1, 0]
+    # np.save('MeshUtils\stretched_uv_map_3.npy', final)
+
+    pass

custom_utils/tps_custom.py

@@ -0,0 +1,210 @@
+import numpy
+import pygeodesic.geodesic as geodesic
+
+__all__ = ['find_coefficients', 'transform']
+
+
+def cdist(K: numpy.ndarray, B: numpy.ndarray) -> numpy.ndarray:
+    """Calculate Euclidean distance between K[i, :] and B[j, :].
+
+    Arguments
+    ---------
+        K : numpy.array
+        B : numpy.array
+    """
+    K = numpy.atleast_2d(K)
+    B = numpy.atleast_2d(B)
+    assert K.ndim == 2
+    assert B.ndim == 2
+
+    K = numpy.expand_dims(K, 1)
+    B = numpy.expand_dims(B, 0)
+    D = K - B
+    return numpy.linalg.norm(D, axis=2)
+
+
+# def cdist(K: numpy.ndarray, B: numpy.ndarray, faces: numpy.ndarray) -> numpy.ndarray:
+#     geoalg = geodesic.PyGeodesicAlgorithmExact(points, faces)
+
+
+
+def pairwise_radial_basis(K: numpy.ndarray, B: numpy.ndarray) -> numpy.ndarray:
+    """Compute the TPS radial basis function phi(r) between every row-pair of K
+    and B where r is the Euclidean distance.
+
+    Arguments
+    ---------
+        K : numpy.array
+            n by d vector containing n d-dimensional points.
+        B : numpy.array
+            m by d vector containing m d-dimensional points.
+
+    Return
+    ------
+        P : numpy.array
+            n by m matrix where.
+            P(i, j) = phi( norm( K(i,:) - B(j,:) ) ),
+            where phi(r) = r^2*log(r), if r >= 1
+                           r*log(r^r), if r <  1
+    """
+    # r_mat(i, j) is the Euclidean distance between K(i, :) and B(j, :).
+    r_mat = cdist(K, B)
+
+    pwise_cond_ind1 = r_mat >= 1
+    pwise_cond_ind2 = r_mat < 1
+    r_mat_p1 = r_mat[pwise_cond_ind1]
+    r_mat_p2 = r_mat[pwise_cond_ind2]
+
+    # P correcponds to the matrix K from [1].
+    P = numpy.empty(r_mat.shape)
+    P[pwise_cond_ind1] = (r_mat_p1**2) * numpy.log(r_mat_p1)
+    P[pwise_cond_ind2] = r_mat_p2 * numpy.log(numpy.power(r_mat_p2, r_mat_p2))
+
+    return P
+
+
+def find_coefficients(control_points: numpy.ndarray,
+                      target_points: numpy.ndarray,
+                      lambda_: float = 0.,
+                      solver: str = 'exact') -> numpy.ndarray:
+    """Given a set of control points and their corresponding points, compute the
+    coefficients of the TPS interpolant deforming surface.
+
+    Arguments
+    ---------
+        control_points : numpy.array
+            p by d vector of control points
+        target_points : numpy.array
+            p by d vector of corresponding target points on the deformed
+            surface
+        lambda_ : float
+            regularization parameter
+        solver : str
+            the solver to get the coefficients. default is 'exact' for the exact
+            solution. Or use 'lstsq' for the least square solution.
+
+    Return
+    ------
+        coef : numpy.ndarray
+            the coefficients
+
+    .. seealso::
+
+        http://cseweb.ucsd.edu/~sjb/pami_tps.pdf
+    """
+    # ensure data type and shape
+    control_points = numpy.atleast_2d(control_points)
+    target_points = numpy.atleast_2d(target_points)
+    if control_points.shape != target_points.shape:
+        raise ValueError(
+            'Shape of and control points {cp} and target points {tp} are not the same.'.
+            format(cp=control_points.shape, tp=target_points.shape))
+
+    p, d = control_points.shape
+
+    # The matrix
+    K = pairwise_radial_basis(control_points, control_points)
+    P = numpy.hstack([numpy.ones((p, 1)), control_points])
+
+    # Relax the exact interpolation requirement by means of regularization.
+    K = K + lambda_ * numpy.identity(p)
+
+    # Target points
+    M = numpy.vstack([
+        numpy.hstack([K, P]),
+        numpy.hstack([P.T, numpy.zeros((d + 1, d + 1))])
+    ])
+    Y = numpy.vstack([target_points, numpy.zeros((d + 1, d))])
+
+    # solve for M*X = Y.
+    # At least d+1 control points should not be in a subspace; e.g. for d=2, at
+    # least 3 points are not on a straight line. Otherwise M will be singular.
+    solver = solver.lower()
+    if solver == 'exact':
+        X = numpy.linalg.solve(M, Y)
+    elif solver == 'lstsq':
+        X, _, _, _ = numpy.linalg.lstsq(M, Y, None)
+    else:
+        raise ValueError('Unknown solver: ' + solver)
+
+    return X
+
+
+def transform(source_points: numpy.ndarray, control_points: numpy.ndarray,
+              coefficient: numpy.ndarray) -> numpy.ndarray:
+    """Transform the source points form the original surface to the destination
+    (deformed) surface.
+
+    Arguments
+    ---------
+        source_points : numpy.array
+            n by d array of source points to be transformed
+        control_points : numpy.array
+            the control points used in the function `find_coefficients`
+        coefficient : numpy.array
+            the computed coefficients
+
+    Return
+    ------
+        deformed_points : numpy.array
+            n by d array of the transformed point on the target surface
+    """
+    source_points = numpy.atleast_2d(source_points)
+    control_points = numpy.atleast_2d(control_points)
+    if source_points.shape[-1] != control_points.shape[-1]:
+        raise ValueError(
+            'Dimension of source points ({sd}D) and control points ({cd}D) are not the same.'.
+            format(sd=source_points.shape[-1], cd=control_points.shape[-1]))
+
+    n = source_points.shape[0]
+
+    A = pairwise_radial_basis(source_points, control_points)
+    K = numpy.hstack([A, numpy.ones((n, 1)), source_points])
+
+    deformed_points = numpy.dot(K, coefficient)
+    return deformed_points
+
+
+
+
+class TPS:
+    """The thin plate spline deformation warpping.
+    """
+
+    def __init__(self,
+                 control_points: numpy.ndarray,
+                 target_points: numpy.ndarray,
+                 lambda_: float = 0.,
+                 solver: str = 'exact'):
+        """Create a instance that preserve the TPS coefficients.
+
+        Arguments
+        ---------
+            control_points : numpy.array
+                p by d vector of control points
+            target_points : numpy.array
+                p by d vector of corresponding target points on the deformed
+                surface
+            lambda_ : float
+                regularization parameter
+            solver : str
+                the solver to get the coefficients. default is 'exact' for the
+                exact solution. Or use 'lstsq' for the least square solution.
+        """
+        self.control_points = control_points
+        self.coefficient = find_coefficients(
+            control_points, target_points, lambda_, solver)
+
+    def __call__(self, source_points):
+        """Transform the source points form the original surface to the
+        destination (deformed) surface.
+
+        Arguments
+        ---------
+            source_points : numpy.array
+                n by d array of source points to be transformed
+        """
+        return transform(source_points, self.control_points,
+                                   self.coefficient)
+
+    transform = __call__

cvae.py

@@ -1,6 +1,6 @@
 import os
 import sys
-from configs.config_abl_norm import config
+from configs.config_abl_all_post import config
 
 os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
 import tensorflow as tf
@@ -41,11 +41,11 @@ from losses import collision_loss as collision_penalty
 from custom_utils.custom_utils import batch_to_dict, load_model, plot_3d, UV2Mesh, get_summary_html, maskFacesBatch, _maskFaces, plot_3d_collision
 from dataset import DatasetFactory, StaticDataset
 from custom_utils.smpl import SMPL
-from models.cvae_model import CVAE
+from models.cvae_model import CVAENew as CVAE
 
 
 def init_wandb():
-    run = wandb.init(project='ablation', name=config['name'], notes=config['notes'], dir='/home/workspace/local-data/temp')
+    run = wandb.init(project='ablation', name=config['name'], notes=config['notes'], dir='/wandb')
     artifact = wandb.Artifact('code', type='code')
     artifact.add_file('cvae.py')
     artifact.add_file('models/cvae_model.py')
@@ -61,7 +61,7 @@ def init_wandb():
         artifact.add_file('/'.join(config_file_name.split('.')) + '.py')
     run.log_artifact(artifact)
 
-uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_adaptive=False)
+uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_adaptive=False, is_smpl_orig=config['is_cape'])
 
 
 def kl_loss(mean, logvar):
@@ -71,7 +71,9 @@ def kl_loss(mean, logvar):
     return kl_loss
 
 
-@tf.function
+# 3d loss on posed !!!!!!!!!!!!
+
+# @tf.function
 def train_step(batch):
 
     with tf.GradientTape() as tape:
@@ -100,14 +102,17 @@ def train_step(batch):
         c_loss, c_ratio = collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
         c_loss, c_ratio = tf.reduce_mean(c_loss), tf.reduce_mean(c_ratio)
         l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((uv_orig - pred) ** 2), axis=-1), axis=[1, 2]))
-        l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1), axis=-1))
+        l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, axis=-1), axis=-1))
-        n_loss = normal_loss(mesh_pred, mesh_orig, batch['mesh_mask'], is_skirt=config['is_skirt'])
+        n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
         # conc_loss = concept_loss(pred, batch['uv_unposed'])
         kl = kl_loss(mean, logvar)
 
         reconstr_loss = l2_loss
         if config['use_extra_losses']:
-            reconstr_loss = reconstr_loss + l2_mesh_loss + (e_loss + c_loss) * config['reg_loss_coeff']
+            if 'block_collision_loss' in config.keys() and config['block_collision_loss'] == True:
+                reconstr_loss = reconstr_loss + l2_mesh_loss + (e_loss) * config['reg_loss_coeff']
+            else:
+                reconstr_loss = reconstr_loss + l2_mesh_loss + (e_loss + c_loss) * config['reg_loss_coeff']
         if config['use_norm_loss'] and step > 2000:
             reconstr_loss = reconstr_loss + n_loss * config['reg_loss_coeff']
 
@@ -155,8 +160,8 @@ def valid_step(batch):
     c_loss, c_ratio = collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
     c_loss, c_ratio = tf.reduce_mean(c_loss), tf.reduce_mean(c_ratio)
     l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((uv_orig - pred) ** 2), axis=-1), axis=[1, 2]))
-    l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1), axis=-1))
+    l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, axis=-1), axis=-1))
-    n_loss = normal_loss(mesh_pred, mesh_orig, batch['mesh_mask'], is_skirt=config['is_skirt'])
+    n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
     # conc_loss = concept_loss(pred, batch['uv_unposed'])
     kl = kl_loss(mean, logvar)
 
@@ -256,7 +261,7 @@ def plot_everything(batch: dict, type: str):
     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" + skirt_modifier + ".npy")
+    F = np.load("custom_utils/smpl_faces" + skirt_modifier + ("_orig" if config['is_cape'] else "") + ".npy")
     mask = batch['mesh_mask'].numpy().astype(bool)
 
     if config['use_skinning']:
@@ -284,7 +289,7 @@ def plot_everything(batch: dict, type: str):
 
     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_faces = [np.concatenate((F, faces[i] + config['smpl_res']), 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))]
 
     garment_posed_fixed = [mesh_posed_fixed[i, mask[i], :] for i in range(body.shape[0])]
@@ -294,8 +299,8 @@ def plot_everything(batch: dict, type: str):
     figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], faces[i], collision[i], collision_faces[i], collision_color[i], flip_xy=(not config['use_skinning'] and config['remove_global_orientation'])) for i in range(4) }
     figs_coll = { type + "/garment_" + str(i) + '_collision': plot_3d_collision(collision_gt[i], collision[i], collision_fixed[i], collision_faces[i], collision_color[i], flip_xy=(not config['use_skinning'] and config['remove_global_orientation'])) for i in range(4) }
 
-    mean_hist = wandb.Histogram(sequence=mean.numpy().mean(axis=0), num_bins=model.latent_dim)
+    mean_hist = wandb.Histogram(sequence=mean.numpy().mean(axis=0))
-    logvar_hist = wandb.Histogram(sequence=logvar.numpy().mean(axis=0), num_bins=model.latent_dim)
+    logvar_hist = wandb.Histogram(sequence=logvar.numpy().mean(axis=0))
 
     return { **figs, **figs_coll, **{ type + '/images': imgs, type + '/mean_hist': mean_hist, type + '/logvar_hist': logvar_hist } }
 
@@ -323,7 +328,7 @@ def sample(folder):
 
         F = np.load("custom_utils/smpl_faces.npy")
         mask = batch['mesh_mask'].numpy().astype(bool)[0]
-        _, faces = _maskFaces(np.ones((14475, 3)), F, mask)
+        _, faces = _maskFaces(np.ones((config['smpl_res'], 3)), F, mask)
 
         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(pred_zero)
@@ -471,12 +476,12 @@ def test():
 
 
 if __name__ == '__main__':
-    init_wandb()
+    # init_wandb()
     parser = argparse.ArgumentParser(description='')
     parser.add_argument('--name', type=str)
     args = vars(parser.parse_args())
 
-    smpl = SMPL(use_one_hot=config['use_one_hot'])
+    smpl = SMPL(orig_smpl=config['is_cape'])
     skirt_modifier = '_skirt' if config['is_skirt'] else ''
 
     factory = DatasetFactory()
@@ -487,11 +492,14 @@ if __name__ == '__main__':
                                         resolutions=[config['resolution']], 
                                         in_memory=config['dataset_in_memory'],
                                         use_one_hot=config['use_one_hot'],
-                                        use_no_orient=config['remove_global_orientation'])
+                                        use_no_orient=config['remove_global_orientation'],
+                                        frames_file=config['frames_file'] if 'frames_file' in config.keys() else None, 
+                                        is_cape=config['is_cape'])
     static_dataset_train = StaticDataset(folder=config['train_dir'],
                                          signature=config['signature_static'], 
                                          resolutions=[config['resolution']], 
-                                         in_memory=True)
+                                         in_memory=True, 
+                                        is_cape=config['is_cape'])
 
     dataset_valid = factory.get_dataset(folder=config['valid_dir'],
                                         batch_size=config['batch_size'], 
@@ -499,11 +507,14 @@ if __name__ == '__main__':
                                         resolutions=[config['resolution']], 
                                         in_memory=config['dataset_in_memory'], 
                                         use_one_hot=config['use_one_hot'],
-                                        use_no_orient=config['remove_global_orientation'])
+                                        use_no_orient=config['remove_global_orientation'], 
+                                        frames_file=config['frames_file_valid'] if 'frames_file_valid' in config.keys() else None, 
+                                        is_cape=config['is_cape'])
     static_dataset_valid = StaticDataset(folder=config['valid_dir'],
                                          signature=config['signature_static'], 
                                          resolutions=[config['resolution']], 
-                                         in_memory=True)
+                                         in_memory=True, 
+                                        is_cape=config['is_cape'])
 
     model = CVAE(config)
 
@@ -515,11 +526,14 @@ if __name__ == '__main__':
     # test()
 
     ckpt = tf.train.Checkpoint(model=model)
-    manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + config['name'], max_to_keep=20)
+    manager = tf.train.CheckpointManager(ckpt, 'checkpoints/post/' + config['name'], max_to_keep=20)
+
 
     if config['pretrained_model'] != '':
         load_model(model, config['pretrained_model'])
 
+    # load_model(model, 'checkpoints/' + config['name'])
+ 
     valid_iterator = iter(dataset_valid)
     step = 0
     for batch in dataset_train:
@@ -537,15 +551,15 @@ if __name__ == '__main__':
             valid_figs = plot_everything(data_valid, 'valid')
             to_log = { **to_log, **train_figs, **valid_figs }
 
-        if step % 10000 == 0:
+        if step % 5000 == 0:
             manager.save(checkpoint_number=step)
         if step == 0:
             to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
 
-        if train_loss['train/euclid_error'] < 0.017:
+        if train_loss['train/euclid_error'] < (config['balancing_activation_limit'] if 'balancing_activation_limit' in config.keys() else 0.017):
             if balancing_activated_step == 0:
                 balancing_activated_step = step
-            g_balancing_coeff = g_balancing_coeff + 1 / 20000      
+            g_balancing_coeff = g_balancing_coeff + 1 / 10000      
 
         to_log = { **to_log, **{ 'train/balancing_coeff': g_balancing_coeff } }
         

cvae_curv.py

@@ -1,5 +1,8 @@
 import os
-os.environ["CUDA_VISIBLE_DEVICES"] = "1"
+import sys
+from configs.config_cvae_curv import config
+
+os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
 import tensorflow as tf
 
 gpus = tf.config.experimental.list_physical_devices('GPU')
@@ -21,9 +24,10 @@ import random
 import cmapy as cm
 import numpy as np
 import scipy as sp
-from custom_utils.custom_utils import UV2Mesh, UV2MeshLabel, maskFacesBatch, _maskFaces
+from custom_utils.custom_utils import UV2Mesh, UV2MeshLabel, maskFacesBatch, _maskFaces, load_model, get_summary_html, get_summary_str
 import plotly.graph_objects as go
 import plotly.express as ex
+from plotly.subplots import make_subplots
 from PIL import Image
 
 from tf_utils import check_tf_version, allow_memory_growth
@@ -33,31 +37,31 @@ from dataset import DatasetFactory, StaticDataset
 from custom_utils.custom_utils import batch_to_dict
 from custom_utils.smpl import SMPL
 from models.cvae_curv_model import CVAE
+from losses import kl_loss
+
 
 
-name='cvae_curv_3'
 def init_wandb():
-    run = wandb.init(project='vae', name=name,
+    run = wandb.init(project='ablation', name=config['name'], notes=config['notes'], dir='/wandb')
-            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.. r lower kl coeff to 0.1. fixed plotting colors\
-addent MLP for lowering conditional Z diemnsion\
-higher starting kl coeff (0.1)\
-train on 100 garments')
     artifact = wandb.Artifact('code', type='code')
     artifact.add_file('cvae_curv.py')
     artifact.add_file('models/cvae_curv_model.py')
     artifact.add_file('dataset.py')
     artifact.add_file('losses.py')
+    artifact.add_file('models/base.py')
+    artifact.add_file('models/cond_encoder.py')
+    artifact.add_file('dataset.py')
     artifact.add_file('convnext/models/convnext_tf.py')
+    config_file_names = [key for key in sys.modules.keys()  if key.startswith('configs.')]
+    for config_file_name in config_file_names:
+        artifact.add_file('/'.join(config_file_name.split('.')) + '.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
 
+def plot_3d(curv_gt, curv_pred, faces, mesh):
     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))
@@ -65,58 +69,50 @@ def plot_3d(curv_gt, curv_pred, faces, mesh):
         v_color_1[curv_gt == i, :] = colors[i]
         v_color_2[curv_pred == i, :] = colors[i]
 
-    fig = go.Figure(data=[
+    fig = make_subplots(rows=1, cols=2, start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(2)] ])
-        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',
+    fig.add_trace(go.Mesh3d(
-                    scene_aspectratio=dict(x=1, y=2, z=1))
+                # 8 vertices of a cube
-    fig.update_layout(
+                x=mesh[:,0],
-        scene = dict(
+                y=mesh[:,1],
-            xaxis = dict(nticks=4, range=[-1,1],),
+                z=mesh[:,2],
-                        yaxis = dict(nticks=4, range=[-1,3],),
+                # i, j and k give the vertices of triangles
-                        zaxis = dict(nticks=4, range=[-1,1],),)
+                i = faces[:,0],
+                j = faces[:,1],
+                k = faces[:,2],
+                vertexcolor=v_color_1,
+                text='gt'
+            ),
+            row=1, col=1)
+    
+    fig.add_trace(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_2,
+                text='pred'
+            ),
+            row=1, col=2)
+
+
+    fig.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],),)
+    fig.update_layout(scene = helper,
+                        scene2 = helper,
                         )
 
     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):
@@ -157,83 +153,89 @@ def dice_dense(prediction, ground_truth, weight_map=None):
 @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']])
+        pred, mean, logvar = model([batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']])
 
-        ce_uv_loss = scce(batch['uv_curvature'], pred)
+        ce_uv_loss = scce(tf.math.argmax(batch['uv_curvature'], axis=-1), pred)
-        ms_ssim = tf.reduce_mean(tf.image.ssim(pred, uv_curvature, max_val=1))
+        dice = dice_dense(pred, batch['uv_curvature'])
 
-        reconstr_loss = ce_uv_loss
+        mesh_pred = uv2mesh(pred) * batch['mesh_mask'][:,:,tf.newaxis]
+        mesh_orig = tf.cast(uv2mesh(tf.math.argmax(batch['uv_curvature'], axis=-1, output_type=tf.int32)), tf.float32) * batch['mesh_mask']
+        ce_mesh_loss = scce(mesh_orig, mesh_pred)
+
+        reconstr_loss = ce_uv_loss + ce_mesh_loss
+        # reconstr_loss = 1 - dice
         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)
+        loss = reconstr_loss + g_balancing_coeff * kl
         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]
+    mean_iou = get_mean_iou(tf.math.argmax(batch['uv_curvature'], axis=-1), tf.math.argmax(pred, axis=-1))
-    # 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
+    return { 'train/loss': loss, 'train/ce_uv_loss': ce_uv_loss, 'train/mean_iou': mean_iou, 'train/dice': dice,
-             'train/reconstr_loss': reconstr_loss,
+             'train/reconstr_loss': reconstr_loss, 'train/ce_mesh_loss': ce_mesh_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/mean': mean, 'train/logvar': logvar, 'train/kl_loss': kl,'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 = model([batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']])
 
-    pred, mean, logvar, z = model([batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask']])
+    ce_uv_loss = scce(tf.math.argmax(batch['uv_curvature'], axis=-1), pred)
-
-    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))
+    mean_iou = get_mean_iou(tf.math.argmax(batch['uv_curvature'], axis=-1), tf.math.argmax(pred, axis=-1))
-    dice = dice_dense(pred, uv_curvature)
+    dice = dice_dense(pred, batch['uv_curvature'])
 
-    reconstr_loss = tf.reduce_mean(ce_uv_loss)
+    reconstr_loss = ce_uv_loss
+    # reconstr_loss = 1 - dice
+    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)
+    loss = reconstr_loss + g_balancing_coeff * kl
     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,
+    mesh_pred = uv2mesh(pred) * batch['mesh_mask'][:,:,tf.newaxis]
-            'valid/ms_ssim': ms_ssim, 'valid/reconstr_loss': reconstr_loss,
+    mesh_orig = tf.cast(uv2mesh(tf.math.argmax(batch['uv_curvature'], axis=-1, output_type=tf.int32)), tf.float32) * batch['mesh_mask']
+    ce_mesh_loss = scce(mesh_orig, mesh_pred)
+
+    return { 'valid/loss': loss, 'valid/ce_uv_loss': ce_uv_loss, 'valid/mean_iou': mean_iou, 'valid/dice': dice,
+             'valid/reconstr_loss': reconstr_loss, 'valid/ce_mesh_loss': ce_mesh_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 = model([batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']])
-    pred, mean, logvar, z = model([batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask']])
     
     pred = np.argmax(pred.numpy(), axis=-1)
     mesh_pred = uv2mesh(pred)
-    mesh_orig = batch['mesh_curvature']
+    mesh_orig = uv2mesh(tf.math.argmax(batch['uv_curvature'], axis=-1)) # 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)
+    img = tf.concat((tf.math.argmax(orig_img, axis=-1), pred_img), axis=2)
+
+    colors = np.array([cm.color('jet', 1 - i) for i in np.linspace(0, 1, 6+1)])
+    color_img = np.zeros((img.shape[0], img.shape[1], img.shape[2], 3), np.uint8)
+    for i in range(7):
+        color_img[img == i, :] = colors[i]
+
     imgs = [wandb.Image(img[i].numpy().astype(np.float32) / 6.0) for i in range(4)]
+    color_img = [wandb.Image(color_img[i]) 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])]
+    gt_curv = [mesh_orig[i].numpy()[mask[i]] 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])]
 
@@ -242,14 +244,9 @@ def plot_everything(batch: dict, type: str):
     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 } }
+    return { **figs, **{ type + '/images': imgs, type + '/images_color': color_img, 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)
@@ -335,101 +332,80 @@ def sample(folder):
 
 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()
+    skirt_modifier = '_skirt' if config['is_skirt'] else ''
     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/train_100', batch_size=batch_size, signature=signature_dynamic, use_filled=False, resolutions=[512])
+    dataset_train = factory.get_dataset(folder=config['train_dir'],
-    static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', signature=signature_static, use_filled=False, resolutions=[512])
+                                        batch_size=config['batch_size'], 
-    dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/valid', batch_size=batch_size, signature=signature_dynamic)
+                                        signature=config['signature_dynamic'],
-    static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/valid', signature=signature_static)
+                                        resolutions=[config['resolution']], 
+                                        in_memory=config['dataset_in_memory'],
+                                        use_one_hot=False,
+                                        use_no_orient=True,
+                                        frames_file=config['frames_file'] if 'frames_file' in config.keys() else None)
+    static_dataset_train = StaticDataset(folder=config['train_dir'],
+                                         signature=config['signature_static'], 
+                                         resolutions=[config['resolution']], 
+                                         in_memory=True)
 
+    dataset_valid = factory.get_dataset(folder=config['valid_dir'],
+                                        batch_size=config['batch_size'], 
+                                        signature=config['signature_dynamic'], 
+                                        resolutions=[config['resolution']], 
+                                        in_memory=config['dataset_in_memory'], 
+                                        use_one_hot=False,
+                                        use_no_orient=True)
+    static_dataset_valid = StaticDataset(folder=config['valid_dir'],
+                                         signature=config['signature_static'], 
+                                         resolutions=[config['resolution']], 
+                                         in_memory=True)
 
-    model = CVAE(latent_dim=128, static_model='checkpoints/ae_static_5',
+    model = CVAE(config)
-                 body_model='checpoints/ae_body_8',)
-
 
     # 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)
+    optimizer = tf.optimizers.Adam(learning_rate=config['lr']) #0.0001
-    balancing_coeff = 0.1
+    g_balancing_coeff = config['starting_balancing_coeff']
-    reg_coeff = 1
     folder = 'models/ae_logs/'
     valid_iterator = iter(dataset_valid)
-    gamma = 1
-    balancing_activated_step = 0
-
-
-    # sample('checkpoints/cvae_curv_2')
-
 
     ckpt = tf.train.Checkpoint(model=model)
-    manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + name, max_to_keep=20)
+    manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + config['name'], max_to_keep=20)
+
+    if config['pretrained_model'] != '':
+        load_model(model, config['pretrained_model'], iteration=config['pretrained_model_iteration'])
 
     step = 0
     for batch in dataset_train:
-        batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic))
+        batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
 
         train_loss = train_step(batch)
 
-        # data_valid = valid_iterator.get_next()
+        data_valid = valid_iterator.get_next()
-        # data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, signature_dynamic))
+        data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, config['signature_dynamic']))
-        # valid_loss = valid_step(data_valid)
+        valid_loss = valid_step(data_valid)
-        valid_loss = {}
 
         to_log = { **train_loss, **valid_loss, **{ 'current_step': step} }
-        if step % 2500 == 0:
+        if step % 5000 == 0 or (step < 5000 and step % 500 == 0):
             train_figs = plot_everything(batch, 'train')
-            # valid_figs = plot_everything(data_valid, 'valid')
+            valid_figs = plot_everything(data_valid, 'valid')
-            valid_figs = {}
             to_log = { **to_log, **train_figs, **valid_figs }
 
-        if step % 25000 == 0:
+        if step % 10000 == 0:
             manager.save(checkpoint_number=step)
+        if step == 0:
+            to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
 
-        if step > 25000:
+        if step > 40000:
-        # if train_loss['train/euclid_error'] < 0.01:
+            g_balancing_coeff = g_balancing_coeff + 1 / 20000
-            balancing_coeff = balancing_coeff + 1 / 20000
-        # if train_loss['train/euclid_error'] < 0.015:
-        #     reg_coeff = reg_coeff + 100 / 15000
 
-
+        to_log = { **to_log, **{ 'train/g_balancing_coeff': g_balancing_coeff } }
-        # 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)
 

cvae_evaluate.py

@@ -0,0 +1,202 @@
+import os
+import string
+from configs.config_cvae import config
+os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
+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, get_summary_html, maskFacesBatch, _maskFaces
+import plotly.graph_objects as go
+import plotly.express as ex
+from PIL import Image
+
+if config['is_skirt']:
+    from losses import edge_loss_skirt as edge_loss
+else:
+    from losses import edge_loss
+
+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 collision_ratio, uv_mask_loss, bend_loss, kl_loss
+# from losses import collision_penalty as collision_penalty
+from losses import collision_loss as collision_penalty
+from custom_utils.custom_utils import batch_to_dict, load_model, plot_3d
+from dataset import DatasetFactory, StaticDataset
+from custom_utils.smpl import SMPL
+from custom_utils.supersmpl_io import writeOBJ
+from models.cvae_model import CVAE
+
+
+def init_wandb():
+    run = wandb.init(project='cvae', name=config['name'], notes=config['notes'])
+    artifact = wandb.Artifact('code', type='code')
+    artifact.add_file('cvae.py')
+    artifact.add_file('cvae_evaluate.py')
+    artifact.add_file('models/cvae_model.py')
+    artifact.add_file('losses.py')
+    artifact.add_file('configs/config_base.py')
+    artifact.add_file('configs/config_cvae.py')
+    artifact.add_file('models/base.py')
+    artifact.add_file('models/cond_encoder.py')
+    artifact.add_file('dataset.py')
+    artifact.add_file('convnext/models/convnext_tf.py')
+    run.log_artifact(artifact)
+
+uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'])
+
+
+@tf.function
+def inference_step(batch):
+    pred = model.sample(batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask'])
+
+    mesh_pred = uv2mesh(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_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)
+
+    # is the output already masked?
+
+    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']))
+    # c_ratio = tf.reduce_mean(collision_ratio(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']))
+    l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed'] - pred) ** 2), axis=-1), axis=[1, 2]))
+
+    ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 5, batch['uv_unposed'] + 5, max_val=10, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
+    euclid_error = 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'])
+
+    return { 'test/l2_loss': l2_loss, 'test/edge_loss': e_loss, 'test/collision_loss': c_loss, 'test/euclid_error': euclid_error, 'test/ms_ssim': ms_ssim,
+ }
+
+
+
+def plot_everything(batch: dict, type: str):
+    pred = model.sample(batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask'])
+    
+    mesh_pred = uv2mesh(pred).numpy()
+    mesh_orig = batch['mesh_unposed']
+    
+    orig_img = batch['uv_unposed']
+    uv_static_img = batch['uv_static']
+    uv_body_posed_img = batch['uv_body_posed']
+    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)]
+    
+    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" + skirt_modifier + ".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))]
+
+    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) }
+    
+    for i in range(4):
+        temp_file = 'temp/' + config['name'] + '_' + batch['id'][i].numpy().decode('UTF-8') + '_' + ''.join(random.choice(string.ascii_lowercase) for _ in range(4))
+        writeOBJ(temp_file + '_garment.obj', garment_posed[i], faces[i].astype(np.int32))
+        writeOBJ(temp_file + '_body.obj', body[i], F[i].astype(np.int32))
+
+
+    return { **figs, **{ type + '/images': imgs } }
+
+
+
+def test():
+    load_model(model, 'checkpoints/' + config['name'].split('_')[0])
+
+    dataset_test = factory.get_dataset(folder=config['test_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=False, do_repeat=False)
+    static_dataset_test = StaticDataset(folder=config['test_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=False)
+   
+    step = 0
+    losses = []
+
+    for batch in dataset_test:
+        batch = static_dataset_test.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
+
+        inference_loss = inference_step(batch)
+        losses.append(inference_loss)
+
+        to_log = { **{ 'current_step': step }, **inference_loss }
+        if step % 100 == 0:
+            figs = plot_everything(batch, 'test')
+            to_log = { **to_log, **figs }
+
+        wandb.log(to_log, step = step)
+
+        if step > 300:
+            break
+        
+        print(step)
+        step += 1
+   
+    # averages = dict()
+    path_to_plotly_html = "temp/plotly_figure.html"
+
+    histograms = dict()
+    for key in inference_loss.keys():
+        accumulated = [losses[i][key].numpy() for i in range(len(losses))]
+        fig = ex.histogram(accumulated, title='Average: ' + str(np.asarray(accumulated).mean()))
+        # fig.write_html(path_to_plotly_html, auto_play = False)
+
+        histograms['all/' + key.split('/')[-1] + '_hist'] = fig # wandb.Html(path_to_plotly_html)
+        # averages['all/' + key.split('/')[-1] + '_avg'] = np.asarray(accumulated).mean()
+
+    
+
+    wandb.log({ **histograms }, step = step)
+    wandb.finish()
+    pass
+
+
+if __name__ == '__main__':
+    # init_wandb()
+    parser = argparse.ArgumentParser(description='')
+    parser.add_argument('--name', type=str)
+    args = vars(parser.parse_args())
+
+    smpl = SMPL()
+    skirt_modifier = '_skirt' if config['is_skirt'] else ''
+
+    factory = DatasetFactory()
+
+    model = CVAE(config)
+
+    g_balancing_coeff = config['starting_balancing_coeff']
+
+    test()

cvae_pyramid.py

@@ -1,6 +1,6 @@
 import os
 import sys
-from configs.config_abl_pyr_64_128_cape_m import config
+from configs.post.config_abl_pyr_64_128_5 import config
 
 os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
 import tensorflow as tf
@@ -35,18 +35,18 @@ if config['is_skirt']:
 else:
     from losses import edge_loss
 
-from losses import uv_mask_loss, bend_loss, kl_loss, collision_ratio
+from losses import uv_mask_loss, bend_loss, kl_loss, collision_ratio, normal_loss
 # from losses import collision_penalty as collision_penalty
 from losses import collision_loss as collision_penalty
 from custom_utils.custom_utils import batch_to_dict, load_model
 from dataset import DatasetFactory, StaticDataset
 from custom_utils.smpl import SMPL
-from models.pyramid_model_optimized import Pyramid
+from models.pyramid_model_optimized import PyramidNew as Pyramid
 from postprocessing import fix_collision
 
 
 def init_wandb():
-    run = wandb.init(project='ablation', name=config['name'], notes=config['notes'], dir='/wandb')
+    run = wandb.init(project='ablation', name=config['name'], notes=config['notes'], dir='/wandb', tags=config['tags'])
 
     artifact = wandb.Artifact('code', type='code')
     artifact.add_file('cvae_pyramid.py')
@@ -59,7 +59,7 @@ def init_wandb():
     artifact.add_file('models/cond_encoder.py')
     artifact.add_file('dataset.py')
     artifact.add_file('convnext/models/convnext_tf.py')
-    config_file_names = [key for key in sys.modules.keys()  if key.startswith('configs.')]
+    config_file_names = [key for key in sys.modules.keys()  if key.startswith('configs.post.')]
     for config_file_name in config_file_names:
         artifact.add_file('/'.join(config_file_name.split('.')) + '.py')
     run.log_artifact(artifact)
@@ -208,7 +208,6 @@ def plot_3d(mesh_gt, mesh_pred_unposed_levels, faces, mesh_posed, collision, col
     return fig
 
 
-
 @tf.function
 def train_step(batch):
     levels_input, conds = group_inputs_for_levels(batch, config)
@@ -251,7 +250,8 @@ def train_step(batch):
         l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((uv_orig - pred) ** 2), axis=-1), axis=[1, 2]))
         # ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 5, batch['uv_unposed'] + 5, max_val=10, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
         # ms_ssim = tf.reduce_mean(tf.image.ssim(pred + 1, batch['uv_unposed'] + 1, max_val=3))
-        l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1), axis=-1))
+        l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, axis=-1), axis=-1))
+        n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_mask'], is_skirt=config['is_skirt'])
         kl = kl_loss(mean, logvar)
 
         if config['use_extra_losses']:
@@ -261,6 +261,8 @@ def train_step(batch):
         else:
             reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (l2_loss)
 
+        if config['use_norm_loss'] and step > 2000:
+            reconstr_loss_total = reconstr_loss_total + n_loss * config['reg_loss_coeff']
         
         kl_total += kl
 
@@ -271,7 +273,7 @@ def train_step(batch):
 
         output = { **output, **{ 'train/edge_loss' + res_modifier: e_loss, 'train/collision_loss' + res_modifier: c_loss,'train/euclid_error' + res_modifier: metric, 
                                  'train/mean' + res_modifier: mean, 'train/logvar' + res_modifier: logvar, 'train/kl_loss' + res_modifier: kl,
-                                 'train/var' + res_modifier: np.e ** logvar, 'train/l2_loss' + res_modifier: l2_loss,
+                                 'train/var' + res_modifier: np.e ** logvar, 'train/l2_loss' + res_modifier: l2_loss, 'train/normal_loss': n_loss,
                                  'train/euclid_error': metric, 'train/l2_mesh_loss': l2_mesh_loss } } # 'train/ms_ssim' + res_modifier: ms_ssim,
 
         loss = reconstr_loss_total + balancing_coeff * kl_total
@@ -321,7 +323,8 @@ def valid_step(batch):
     l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((uv_orig - pred) ** 2), axis=-1), axis=[1, 2]))
     # ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 5, batch['uv_unposed'] + 5, max_val=10, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
     # ms_ssim = tf.reduce_mean(tf.image.ssim(pred + 1, batch['uv_unposed'] + 1, max_val=3))
-    l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1), axis=-1))
+    l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, axis=-1), axis=-1))
+    n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_mask'], is_skirt=config['is_skirt'])
     kl = kl_loss(mean, logvar)
 
     if config['use_extra_losses']:
@@ -331,6 +334,9 @@ def valid_step(batch):
     else:
         reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (l2_loss)
     
+    if config['use_norm_loss'] and step > 2000:
+        reconstr_loss_total = reconstr_loss_total + n_loss * config['reg_loss_coeff']
+        
     kl_total += kl
 
     mean = tf.reduce_mean(mean)
@@ -340,7 +346,7 @@ def valid_step(batch):
 
     output = { **output, **{ 'valid/edge_loss' + res_modifier: e_loss, 'valid/collision_loss' + res_modifier: c_loss,'valid/euclid_error' + res_modifier: metric, 
                                  'valid/mean' + res_modifier: mean, 'valid/logvar' + res_modifier: logvar, 'valid/kl_loss' + res_modifier: kl,
-                                'valid/var' + res_modifier: np.e ** logvar, 'valid/l2_loss' + res_modifier: l2_loss,
+                                'valid/var' + res_modifier: np.e ** logvar, 'valid/l2_loss' + res_modifier: l2_loss, 'valid/normal_loss': n_loss,
                                 'valid/euclid_error': metric, 'valid/l2_mesh_loss': l2_mesh_loss } } # 'valid/ms_ssim' + res_modifier: ms_ssim,
 
     loss = reconstr_loss_total + balancing_coeff * kl_total
@@ -669,13 +675,73 @@ def test():
     pass
 
 
+def generate_results():
+    load_model(model, 'checkpoints/' + config['name'])
+
+    dataset_test = factory.get_dataset(folder=config['test_dir'], 
+                                        batch_size=config['batch_size'], 
+                                        # batch_size=1, 
+                                        signature=config['signature_dynamic'], 
+                                        resolutions=dataset_resolutions, 
+                                        in_memory=config['dataset_in_memory'], 
+                                        is_cape=config['is_cape'],
+                                        frames_file=config['frames_file'] if 'frames_file' in config.keys() else None,
+                                        do_repeat=False,
+                                        shuffle=False)
+    static_dataset_test = StaticDataset(folder=config['test_dir'], 
+                                        signature=config['signature_static'], 
+                                        resolutions=dataset_resolutions, 
+                                        in_memory=True,
+                                        is_cape=config['is_cape'])
+    folder = os.path.join('/home/workspace/data/results/cape', config['name'])
+    os.makedirs(folder, exist_ok=True)
+   
+    for batch in tqdm.tqdm(dataset_test):
+        batch = static_dataset_test.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
+
+        levels_input, conds = group_inputs_for_levels(batch, config)
+        out_levels = model([levels_input, conds], training=False)
+
+        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])])
+
+        i = -1
+        pred = out_levels[i][0]
+
+        mesh_pred = uv2mesh(pred)
+        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)
+
+        for i in range(batch['id'].shape[0]):
+            mask = batch['mesh_mask'][i].numpy().astype(bool)
+            id, frame_nr = batch['id'][i].numpy().decode("utf-8"), batch['frame_nr'][i].numpy().decode("utf-8")
+            unposed = mesh_pred[i].numpy()[mask, :]
+            posed = mesh_pred_posed[i].numpy()[mask, :]
+            np.save(os.path.join(folder, id + '_' + frame_nr + '_posed.npy'), posed)
+            np.save(os.path.join(folder, id + '_' + frame_nr + '_unposed.npy'), unposed)
+
+    pass
+
+@tf.function
+def speed_step(levels_input, conds):
+
+    return model.sample([levels_input, conds], training=False)
+
+
 
 if __name__ == '__main__':
     init_wandb()
     parser = argparse.ArgumentParser(description='')
-    parser.add_argument('--name', type=str)
+    parser.add_argument('--name', type=str, default='lhunor')
     args = vars(parser.parse_args())
 
+    # config_file_names = [key for key in sys.modules.keys()  if key.startswith('configs.')]
+    # assert len(config_file_names) == 1
+    # config_file_name = config_file_names[0]
+    # # assert config_file_name.split('.')[1].endswith(args.name)
+    # assert config_file_name.split('.')[1].endswith(config['name'])
+
     smpl = SMPL(orig_smpl=config['is_cape'])
     skirt_modifier = '_skirt' if config['is_skirt'] else ''
     dataset_resolutions = config['pyramid']['resolutions'].copy()
@@ -683,10 +749,31 @@ if __name__ == '__main__':
         dataset_resolutions.append(512)
 
     factory = DatasetFactory()
-    dataset_train = factory.get_dataset(folder=config['train_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], resolutions=dataset_resolutions, in_memory=config['dataset_in_memory'], is_cape=config['is_cape'])
+    dataset_train = factory.get_dataset(folder=config['train_dir'], 
-    static_dataset_train = StaticDataset(folder=config['train_dir'], signature=config['signature_static'], resolutions=dataset_resolutions, is_cape=config['is_cape'])
+                                        batch_size=config['batch_size'], 
-    dataset_valid = factory.get_dataset(folder=config['valid_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], resolutions=dataset_resolutions, in_memory=config['dataset_in_memory'], is_cape=config['is_cape'])
+                                        # batch_size=1, 
-    static_dataset_valid = StaticDataset(folder=config['valid_dir'], signature=config['signature_static'], resolutions=dataset_resolutions, is_cape=config['is_cape'])
+                                        signature=config['signature_dynamic'], 
+                                        resolutions=dataset_resolutions, 
+                                        in_memory=config['dataset_in_memory'], 
+                                        is_cape=config['is_cape'],
+                                        frames_file=config['frames_file'] if 'frames_file' in config.keys() else None)
+    static_dataset_train = StaticDataset(folder=config['train_dir'], 
+                                        signature=config['signature_static'], 
+                                        resolutions=dataset_resolutions, 
+                                        in_memory=True,
+                                        is_cape=config['is_cape'])
+    dataset_valid = factory.get_dataset(folder=config['valid_dir'], 
+                                        batch_size=config['batch_size'], 
+                                        signature=config['signature_dynamic'], 
+                                        resolutions=dataset_resolutions, 
+                                        in_memory=config['dataset_in_memory'], 
+                                        is_cape=config['is_cape'],
+                                        frames_file=config['frames_file_valid'] if 'frames_file_valid' in config.keys() else None)
+    static_dataset_valid = StaticDataset(folder=config['valid_dir'], 
+                                        signature=config['signature_static'], 
+                                        resolutions=dataset_resolutions, 
+                                        in_memory=True, 
+                                        is_cape=config['is_cape'])
 
     model: Pyramid = Pyramid(config)
 
@@ -700,34 +787,38 @@ if __name__ == '__main__':
     # sample('checkpoints/cvae_13')
     # test()
 
+    # generate_results()
+
     ckpt = tf.train.Checkpoint(model=model)
-    manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + config['name'], max_to_keep=20)
+    manager = tf.train.CheckpointManager(ckpt, 'checkpoints/post/' + config['name'], max_to_keep=20)
 
     # load_model(model, 'checkpoints/pyr_abl_32_skirt', iteration='180000')
 
+    if config['pretrained_model'] != '':
+        # assert len(config['pyramid']['resolutions']) == 1
+        load_model(model, config['pretrained_model'], iteration=config['pretrained_model_iteration'])
+
     if len(config['pyramid']['resolutions']) == 1:
         model.static_enc.trainable = False
         model.body_enc.trainable = False
         model.curv_enc.trainable = False
     else:
-        load_model(model, config['pyramid']['pretrained_pyramid'])
+        load_model(model, config['pyramid']['pretrained_pyramid'], iteration=(config['pyramid']['pretrained_pyramid_iteration'] if 'pretrained_pyramid_iteration' in config.keys() else ''))
         model.static_enc.trainable = False
         model.body_enc.trainable = False
         model.curv_enc.trainable = False
-        model.cond_encoder.trainable = False
 
 
     # load_model(model.levels[0].upscale_enhancer, config['pyramid']['levels'][0]['upscale_enhancer'])
     # model.levels[0].upscale_enhancer.trainable = False
 
-    # the last level doesn't have an upscale enhancer as it's enhanced for the final resolutin directly
+
-    # might not be necessary, as enhancing to final or a bit higher resolution seems to be the same
+    for i in range(len(model.levels)):
-    for i in range(len(model.levels)-1):
         load_model(model.levels[i].upscale_enhancer, config['pyramid']['levels'][i]['upscale_enhancer'])
         model.levels[i].upscale_enhancer.trainable = False
-    if config['pyramid']['resolutions'][-1] != 512:
+    # if config['pyramid']['resolutions'][-1] != 512:
-        load_model(model.upscale_enhancer, config['pyramid']['final_upscale_enhancer'])
+    #     load_model(model.upscale_enhancer, config['pyramid']['final_upscale_enhancer'])
-        model.upscale_enhancer.trainable = False
+    #     model.upscale_enhancer.trainable = False
 
     for i in range(len(model.levels)-1):
         model.levels[i].trainable = False
@@ -738,9 +829,22 @@ if __name__ == '__main__':
 
 
     step = 0
+    times = []
     for batch in dataset_train:
         batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
 
+        # levels_input, conds = group_inputs_for_levels(batch, config)
+
+        # start = time.perf_counter()
+        # out = speed_step(levels_input, conds)
+        # t = time.perf_counter() - start
+        # print(t)
+        # times.append(t)
+
+        # if step > 200:
+        #     break
+
+
         train_loss = train_step(batch)
 
         data_valid = valid_iterator.get_next()
@@ -753,7 +857,7 @@ if __name__ == '__main__':
             valid_figs = plot_everything(data_valid, 'valid')
             to_log = { **to_log, **train_figs, **valid_figs }
 
-        if step % 10000 == 0:
+        if step % 5000 == 0:
             manager.save(checkpoint_number=step)
         if step == 0:
             to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
@@ -771,3 +875,5 @@ if __name__ == '__main__':
         print(step)
 
         step += 1
+    # print('mean:')
+    # print(np.mean(times[1:]))

cvae_pyramid_old.py

@@ -420,7 +420,7 @@ if __name__ == '__main__':
     smpl = SMPL()
 
     factory = DatasetFactory()
-    signature_dynamic = { 'id': (),
+    signature_dynamic = { 'id': (), 'frame_nr': (),
                           'pose_shape': (72 + 11,),
                         #   'uv_posed': (512, 256, 3),
                           'uv_body_posed': (512, 256, 3),

cvae_vae.py

@@ -449,7 +449,7 @@ if __name__ == '__main__':
     smpl = SMPL()
 
     factory = DatasetFactory()
-    signature_dynamic = { 'id': (),
+    signature_dynamic = { 'id': (), 'frame_nr': (),
                           'pose_shape': (72 + 11,),
                         #   'uv_posed': (3, 512, 256,),
                           'uv_body_posed': (3, 512, 256,),

dataset.py

@@ -135,14 +135,27 @@ class StaticDataset():
 
 class DatasetIdx(tf.data.Dataset):
 
-    def __new__(cls, folder, shuffle=True, is_cape=False):
+    def __new__(cls, folder, shuffle=True, is_cape=False, testing=False, frames_file=None):
-        uv_static_folder = os.path.join(folder, 'pose_shape') 
+        if frames_file is None:
-        frames = sorted(os.listdir(uv_static_folder))
+            uv_static_folder = os.path.join(folder, 'pose_shape') 
+            frames = sorted(os.listdir(uv_static_folder))
+
+            if is_cape:
+                frames = ['.'.join(frame.split('.')[:-1]) for frame in frames] # .split('_')[0]
+            else:
+                frames = [frame.split('.')[0] for frame in frames] # .split('_')[0]
+
+            if testing:
+                frames = sorted(os.listdir(os.path.join(folder, 'mesh_static') ))
+                frames = [frame.split('.')[0] + '_4' for frame in frames]
+
+        if frames_file is not None:
+            frames = []
+            with open(os.path.join(folder, frames_file), 'r') as f:
+                lines = f.readlines()
+                lines = [line.strip() for line in lines]
+            frames = lines
 
-        if is_cape:
-            frames = ['.'.join(frame.split('.')[:-1]) for frame in frames] # .split('_')[0]
-        else:
-            frames = [frame.split('.')[0] for frame in frames] # .split('_')[0]
         if shuffle:
             random.shuffle(frames)
 
@@ -179,10 +192,12 @@ class DatasetFactory():
                     do_repeat=True, 
                     is_cape=False, 
                     use_one_hot=False,
-                    use_no_orient=True):
+                    use_no_orient=True,
+                    testing=False,
+                    frames_file=None):
         assert set(resolutions).issubset(set([32, 64, 128, 256, 512]))
 
-        dataset_idx = DatasetIdx(folder, shuffle, is_cape=is_cape)
+        dataset_idx = DatasetIdx(folder, shuffle, is_cape=is_cape, testing=testing, frames_file=frames_file)
 
         separator = '_'
         if is_cape:
@@ -333,11 +348,11 @@ if __name__ == '__main__':
     os.environ["CUDA_VISIBLE_DEVICES"] = "6"
     from custom_utils.custom_utils import batch_to_dict, UV2Mesh, _maskFaces
     import plotly.graph_objects as go
-    # import plotly.express as ex
+    import plotly.express as ex
     # import cmapy as cm
     uv2mesh = UV2Mesh(is_skirt=True)
 
-    F = np.load("custom_utils/smpl_faces_skirt.npy")
+    F = np.load("custom_utils/smpl_faces.npy")
 
 
     # dataset = DatasetIdx('/home/workspace/data/dataset/stylegan_new/train')
@@ -355,7 +370,7 @@ if __name__ == '__main__':
                         #   'mesh_unposed': (14475, 3,),
                           'uv_unposed': (512, 256, 3),
                         #   'uv_unposed_32': (3, 32, 16,),
-                        #   'uv_curvature': (512, 256, 7),
+                          'uv_curvature': (512, 256, 7),
                         #   'mesh_curvature':  (14475)
                         }
     signature_static = { #'uv_static': (3, 512, 256,),
@@ -367,14 +382,17 @@ if __name__ == '__main__':
                           }
 
     dataset = factory.get_dataset(folder='/home/workspace/local-data/train_skirt',
-                                  batch_size=300,
+                                  batch_size=1,
                                   signature=signature_dynamic,
                                   resolutions=[512],
                                   do_repeat=False,
-                                  use_one_hot=False)
+                                  use_one_hot=False,
+                                  is_cape=False,
+                                  frames_file='frames_dress_30k.txt')
     static_dataset = StaticDataset(folder='/home/workspace/local-data/train_skirt',
                                    signature=signature_static,
-                                   resolutions=[512])
+                                   resolutions=[512],
+                                  is_cape=False)
 
 
     # for i, batch in enumerate(dataset):
@@ -386,6 +404,7 @@ if __name__ == '__main__':
     for i, batch in enumerate(dataset):
         print(i)
         batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic))
+        ex.imshow(batch['uv_curvature'][0].numpy().argmax(-1) / 6.0).show()
         sum_uv += batch['uv_unposed'].numpy().sum(axis=0)
         sum_mask += batch['uv_mask'].numpy().sum(axis=0)
 
@@ -406,9 +425,9 @@ if __name__ == '__main__':
                 k = F[:,2]
             )]).show()
 
-    np.save('custom_data/train_no_shape_skirt_avg_uv.npy', avg_garment)
+    np.save('custom_data/train_dress_30k_no_shape_avg_uv.npy', avg_garment)
 
-    # print()
+    print()
 
     # pass
 

dataset_packed.py

@@ -217,7 +217,7 @@ if __name__ == '__main__':
     #     print(d)
 
     factory = DatasetFactory()
-    signature_dynamic = { 'id': (),
+    signature_dynamic = { 'id': (), 'frame_nr': (),
                           'pose_shape': (72 + 11,),
                         #   'uv_posed': (3, 512, 256,),
                           'uv_body_posed': (512, 256, 3),

deorient copy.py

@@ -0,0 +1,104 @@
+
+
+import os
+import sys
+
+
+os.environ["CUDA_VISIBLE_DEVICES"] = '6'
+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 glob
+import tqdm
+import copy
+import numpy as np
+from custom_utils.smpl import SMPL
+
+# from custom_utils.custom_utils import UV2Mesh, _maskFaces
+import plotly.graph_objects as go
+# import plotly.express as ex
+from scipy.spatial.transform import Rotation as R
+
+
+if __name__ == '__main__':
+    folder = '/home/workspace/local-data/test_skirt_ready'
+    # folder = '/home/workspace/data/dataset/test_ready'
+    # uv2mesh = UV2Mesh(is_skirt=False)
+    smpl = SMPL(use_one_hot=False)
+
+    F = np.load("custom_utils/smpl_faces.npy")
+
+    ids = os.listdir(os.path.join(folder, 'pose_shape'))
+    ids = [id.split('_')[0] for id in ids]
+    ids = set(ids)
+    os.makedirs(os.path.join(folder, 'mesh_posed_no_orient'), exist_ok=True)
+
+    for id in tqdm.tqdm(ids):
+
+        frames = glob.glob(os.path.join(folder, 'pose_shape', id + '_*'))
+        frames = len(frames)
+
+        mask = np.load(os.path.join(folder, 'mesh_mask', id + '.npy')).astype(bool)
+        mesh_posed = [np.load(os.path.join(folder, 'mesh_posed',  id + '_' + str(i) + '.npz')) for i in range(frames)]
+        mesh_posed = [mesh_posed[i]['mesh_posed'] for i in range(frames)]
+        # mesh_unposed = [np.load(os.path.join(folder, 'mesh_unposed',  id + '_' + str(i) + '.npz')) for i in range(frames)]
+        # mesh_unposed = [mesh_unposed[i]['mesh_unposed'] for i in range(frames)]
+        pose_shape = [np.load(os.path.join(folder, 'pose_shape',  id + '_' + str(i) + '.npy')).astype(np.float32) for i in range(frames)]
+        # pose_shape_no_orient = copy.deepcopy(pose_shape)
+        # for i in range(len(pose_shape)):
+        #     pose_shape_no_orient[i][:3] = 0
+
+        rodr = smpl.run_rodrigues(np.stack([x[:72] for x in pose_shape]).astype(np.float32)).numpy()
+
+        for frame in range(frames):
+            r = R.from_matrix(rodr[frame, 0, :, :])
+            # assert r.as_matrix() == rodr[frame, 0, :, :]
+            out = r.apply(mesh_posed[frame], inverse=True)
+
+            # x = smpl.run_garment(pose_shape_no_orient[frame][:72][np.newaxis,:], pose_shape_no_orient[frame][72:][np.newaxis,:], mesh_unposed[frame][np.newaxis,:,:])[0].numpy()[mask, :]
+            # y = out[mask, :]
+            # _, faces = _maskFaces(np.zeros((14475,3), np.float32), F, mask)
+            # fig = go.Figure(data=[
+            #         go.Mesh3d(
+            #             x=x[:,0],
+            #             y=x[:,1],
+            #             z=x[:,2],
+            #             i = faces[:,0],
+            #             j = faces[:,1],
+            #             k = faces[:,2],
+            #         ),
+            #         go.Mesh3d(
+            #             x=y[:,0],
+            #             y=y[:,1],
+            #             z=y[:,2],
+            #             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()
+
+
+            np.savez_compressed(os.path.join(folder, 'mesh_posed_no_orient', id + '_' + str(frame) + '.npz'), mesh_posed_no_orient=out)

deorient_cape.py

@@ -0,0 +1,114 @@
+
+
+import os
+import sys
+
+
+os.environ["CUDA_VISIBLE_DEVICES"] = '6'
+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 glob
+import tqdm
+import copy
+import numpy as np
+from custom_utils.smpl import SMPL
+
+# from custom_utils.custom_utils import UV2Mesh, _maskFaces
+import plotly.graph_objects as go
+# import plotly.express as ex
+from scipy.spatial.transform import Rotation as R
+
+
+if __name__ == '__main__':
+    folder = '/home/workspace/local-data/cape_processed/train/female'
+    # folder = '/home/workspace/data/dataset/test_ready'
+    # uv2mesh = UV2Mesh(is_skirt=False)
+    
+    smpl = SMPL(use_one_hot=False, orig_smpl=True)
+
+
+    F = np.load("custom_utils/smpl_faces_orig.npy")
+
+    ids = os.listdir(os.path.join(folder, 'pose_shape'))
+    ids = [id.split('.')[0] for id in ids]
+    ids = set(ids)
+    os.makedirs(os.path.join(folder, 'mesh_posed_no_orient'), exist_ok=True)
+
+    for id in tqdm.tqdm(ids):
+
+        frames = sorted(glob.glob(os.path.join(folder, 'pose_shape', id + '.*')))
+        frames = ['.'.join(frame.split('/')[-1].split('.')[:-1]) + '.npz' for frame in frames]
+
+
+        mesh_posed = [np.load(os.path.join(folder, 'mesh_posed',  frame)) for frame in frames]
+        mesh_posed = [mesh_posed[i]['mesh_posed'] for i in range(len(frames))]
+        pose_shape = [np.load(os.path.join(folder, 'pose_shape',  frame[:-1] + 'y')) for frame in frames]
+
+        # mesh_unposed = [np.load(os.path.join(folder, 'mesh_unposed_no_shape',  frame)) for frame in frames]
+        # mesh_unposed = [mesh_unposed[i]['mesh_unposed_no_shape'] for i in range(len(frames))]
+        # shape = np.asarray(pose_shape)[:, 72:].astype(np.float32)
+        # 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_unposed_shaped = [mesh_unposed[i] + shape_offsets.numpy()[i] for i in range(shape.shape[0])]
+
+
+        pose_shape_no_orient = copy.deepcopy(pose_shape)
+        for i in range(len(pose_shape)):
+            pose_shape_no_orient[i][:3] = 0
+
+        rodr = smpl.run_rodrigues(np.stack([x[:72] for x in pose_shape]).astype(np.float32)).numpy()
+
+        for i, frame in enumerate(frames):
+            r = R.from_matrix(rodr[i, 0, :, :])
+            # assert r.as_matrix() == rodr[frame, 0, :, :]
+            out = r.apply(mesh_posed[i], inverse=True)
+
+            # if i == 50:
+            #     x = smpl.run_garment(pose_shape_no_orient[i][:72][np.newaxis,:].astype(np.float32), pose_shape_no_orient[i][72:][np.newaxis,:].astype(np.float32), mesh_unposed_shaped[i][np.newaxis,:,:].astype(np.float32))[0].numpy()
+            #     y = out
+            #     faces = F
+            #     fig = go.Figure(data=[
+            #             go.Mesh3d(
+            #                 x=x[:,0],
+            #                 y=x[:,1],
+            #                 z=x[:,2],
+            #                 i = faces[:,0],
+            #                 j = faces[:,1],
+            #                 k = faces[:,2],
+            #             ),
+            #             go.Mesh3d(
+            #                 x=y[:,0],
+            #                 y=y[:,1],
+            #                 z=y[:,2],
+            #                 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()
+
+
+            np.savez_compressed(os.path.join(folder, 'mesh_posed_no_orient', frame), mesh_posed_no_orient=out)

deshape.py

@@ -26,7 +26,7 @@ if __name__ == '__main__':
     F = np.load("custom_utils/smpl_faces.npy")
 
     factory = DatasetFactory()
-    signature_dynamic = { 'id': (),
+    signature_dynamic = { 'id': (), 'frame_nr': (),
                           'pose_shape': (72 + 11,),
                           'mesh_posed': (14475, 3,),
                           'mesh_unposed': (14475, 3,),

downscale.py

@@ -8,7 +8,7 @@ import numpy as np
 import plotly.graph_objects as go
 import plotly.express as ex
 
-os.environ["CUDA_VISIBLE_DEVICES"] = "3"
+os.environ["CUDA_VISIBLE_DEVICES"] = "6"
 
 gpus = tf.config.experimental.list_physical_devices('GPU')
 if gpus:
@@ -98,6 +98,23 @@ def downscale_uv_maps_posed_body(folder, type, sizes):
             np.savez_compressed(os.path.join(folder, type + '_2_' + str(size), frame), uv_posed_body_no_orient = down)
 
 
+def downscale_uv_maps_posed_no_orient(folder, type, sizes):
+    for size in sizes:
+        os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True)
+
+    frames = os.listdir(os.path.join(folder, type))[int(args.start):int(args.end)]
+
+    for frame in tqdm.tqdm(frames):
+        id, _ = frame.split('_')
+        img = np.load(os.path.join(folder, type, frame)); img = img['uv_posed_no_orient']
+        img = fill(img, uv_masks_body[512])
+        for size in sizes:
+            mask = uv_masks_body[size]
+            down = tf.image.resize(img, size=(size, int(size/2)), method=tf.image.ResizeMethod.LANCZOS5)
+            down = down * mask[:,:,np.newaxis]
+            np.savez_compressed(os.path.join(folder, type + '_' + str(size), frame), uv_posed_no_orient = down)
+
+
 def downscale_uv_maps_curv(folder, type, sizes):
     for size in sizes:
         os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True)
@@ -171,12 +188,13 @@ def downscale_uv_maps_mask(folder, type, size):
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
     parser.add_argument("-f", "--folder", help = "end of range", default='/home/workspace/local-data/train')
-    parser.add_argument("-s", "--start", help = "start of range", default=280000)
+    parser.add_argument("-s", "--start", help = "start of range", default=30700)
-    parser.add_argument("-e", "--end", help = "end of range", default=10000000)
+    parser.add_argument("-e", "--end", help = "end of range", default=45000)
     parser.add_argument("-n", "--name", help = "end of range", default='lhunor')
     args = parser.parse_args()
 
-    sizes = [32, 64, 128, 256]
+    # sizes = [32, 64, 128, 256]
+    sizes = [32]
     folder = args.folder
 
     is_skirt = True
@@ -200,8 +218,9 @@ if __name__ == '__main__':
     # downscale_uv_maps_static(folder=folder, type='uv_static', sizes=sizes)
 
     # downscale_uv_maps_unposed(folder=folder, type='uv_unposed_no_shape', sizes=sizes)
+    downscale_uv_maps_posed_no_orient(folder=folder, type='uv_posed_no_orient', sizes=sizes)
     # downscale_uv_maps_curv(folder=folder, type='uv_curvature', sizes=sizes)
-    downscale_uv_maps_posed_body(folder=folder, type='uv_posed_body_no_orient', sizes=sizes)
+    # downscale_uv_maps_posed_body(folder=folder, type='uv_posed_body_no_orient', sizes=sizes)
 
 
     print()

downscale_cape.py

@@ -26,7 +26,7 @@ def get_img_idxs():
     img_idxs = dict()
     img_coords = dict()
     uv_masks = dict()
-    uv_mask = np.load('custom_utils/mask_256x512.npy')
+    uv_mask = np.load('custom_data/garment_uv_mask_full.npy')
     coords = np.transpose(np.stack(np.nonzero((1 - uv_mask))), (1,0))
     img_idx = np.transpose(np.stack(np.nonzero((uv_mask))), (1,0))
     img_idxs[512] = img_idx
@@ -41,7 +41,17 @@ def get_img_idxs():
         img_coords[res] = coords
         uv_masks[res] = down.astype(np.int32)
 
-    return img_idxs, img_coords, uv_masks
+
+    uv_masks_body = dict()
+    uv_mask_body = np.load('custom_utils/mask_256x512.npy')
+    uv_mask_body = cv2.dilate(uv_mask_body.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2)
+    uv_masks_body[512] = uv_mask_body
+    for res in [32, 64, 128, 256]:
+        down = tf.image.resize(uv_mask_body[:,:,tf.newaxis].astype(np.int32), size=(res, int(res/2)), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
+        down = np.squeeze(down.numpy().astype(bool))
+        uv_masks_body[res] = down
+
+    return img_idxs, img_coords, uv_masks, uv_masks_body
 
 
 def fill(img, mask):
@@ -96,12 +106,10 @@ def downscale_uv_maps_posed_body(folder, type, sizes):
     frames = os.listdir(os.path.join(folder, type))[int(args.start):int(args.end)]
 
     for frame in tqdm.tqdm(frames):
-        id, _ = frame.split('_')
         img = np.load(os.path.join(folder, type, frame)); img = img['uv_posed_body_no_orient']
-        mask_full = uv_mask_store[512][id]
+        img = fill(img, uv_masks_body[512])
-        img = fill(img, mask_full)
         for size in sizes:
-            mask = uv_mask_store[size][id]
+            mask = uv_masks_body[size]
             down = tf.image.resize(img, size=(size, int(size/2)), method=tf.image.ResizeMethod.LANCZOS5)
             down = down * mask[:,:,np.newaxis]
             np.savez_compressed(os.path.join(folder, type + '_' + str(size), frame), uv_posed_body_no_orient = down)
@@ -114,12 +122,11 @@ def downscale_uv_maps_curv(folder, type, sizes):
     frames = os.listdir(os.path.join(folder, type))[int(args.start):int(args.end)]
 
     for frame in tqdm.tqdm(frames):
-        id, _ = frame.split('_')
         img = np.load(os.path.join(folder, type, frame)); img = img['curvature_uv']
-        mask_full = uv_mask_store[512][id]
+        mask_full = uv_mask_store[512][None]
         img = fill_label(img, mask_full)
         for size in sizes:
-            mask = uv_mask_store[size][id]
+            mask = uv_mask_store[size][None]
             down = tf.image.resize(img[:,:,tf.newaxis], size=(size, int(size/2)), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
             down = np.squeeze(down) * mask[:,:]
             np.savez_compressed(os.path.join(folder, type + '_' + str(size), frame), uv_curvature = down)
@@ -133,12 +140,11 @@ def downscale_uv_maps_unposed(folder, type, sizes):
     frames = os.listdir(os.path.join(folder, type))[int(args.start):int(args.end)]
 
     for frame in tqdm.tqdm(frames):
-        id, _ = frame.split('_')
         img = np.load(os.path.join(folder, type, frame)); img = img['uv_unposed_no_shape']
-        mask_full = uv_mask_store[512][id]
+        mask_full = uv_mask_store[512][None]
         img = fill(img, mask_full)
         for size in sizes:
-            mask = uv_mask_store[size][id]
+            mask = uv_mask_store[size][None]
             down = tf.image.resize(img, size=(size, int(size/2)), method=tf.image.ResizeMethod.LANCZOS5)
             down = down * mask[:,:,np.newaxis]
             np.savez_compressed(os.path.join(folder, type + '_' + str(size), frame), uv_unposed_no_shape = down)
@@ -151,14 +157,13 @@ def downscale_uv_maps_static(folder, type, sizes):
     frames = os.listdir(os.path.join(folder, type))[int(args.start):int(args.end)]
 
     for frame in tqdm.tqdm(frames):
-        id, _ = frame.split('.')
         img = np.load(os.path.join(folder, type, frame)); img = img['uv_static']
-        mask_full = uv_mask_store[512][id]
+        mask_full = uv_mask_store[512][None]
         img = fill(img, mask_full)
 
         for size in sizes:
             down = tf.image.resize(img, size=(size, int(size/2)), method=tf.image.ResizeMethod.LANCZOS5)
-            mask = uv_mask_store[size][id]
+            mask = uv_mask_store[size][None]
             down = down * mask[:,:,np.newaxis]
             np.savez_compressed(os.path.join(folder, type + '_' + str(size), frame), uv_static = down)
 
@@ -169,7 +174,6 @@ def downscale_uv_maps_mask(folder, type, size):
     frames = os.listdir(os.path.join(folder, type))[int(args.start):int(args.end)]
 
     for frame in tqdm.tqdm(frames):
-        id, _ = frame.split('.')
         img = np.load(os.path.join(folder, type, frame)); img = img['uv_mask'].astype(np.float32)
         down = tf.image.resize(img[:,:,tf.newaxis], size=(size, int(size/2)), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
         down = np.squeeze(down.numpy().astype(bool))
@@ -179,9 +183,9 @@ def downscale_uv_maps_mask(folder, type, size):
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
-    parser.add_argument("-f", "--folder", help = "end of range", default='/home/workspace/local-data/cape_processed/test/male')
+    parser.add_argument("-f", "--folder", help = "end of range", default='/home/workspace/local-data/cape_processed/train/male')
     parser.add_argument("-s", "--start", help = "start of range", default=0)
-    parser.add_argument("-e", "--end", help = "end of range", default=10000)
+    parser.add_argument("-e", "--end", help = "end of range", default=100000)
     parser.add_argument("-n", "--name", help = "end of range", default='lhunor')
     args = parser.parse_args()
 
@@ -189,7 +193,7 @@ if __name__ == '__main__':
     sizes = [32, 64, 128, 256]
     folder = args.folder
 
-    img_idxs, img_coords, uv_masks = get_img_idxs()
+    img_idxs, img_coords, uv_masks, uv_masks_body = get_img_idxs()
 
     uv_mask_store = { size: UVMaskStore(folder, size=size) for size in sizes }
     uv_mask_store[512] = UVMaskStore(folder, size=512)
@@ -200,9 +204,9 @@ if __name__ == '__main__':
     # downscale_uv_maps_mask(folder=folder, type='uv_mask', size=256)
     downscale_uv_maps_static(folder=folder, type='uv_static', sizes=sizes)
 
-    # downscale_uv_maps_unposed(folder=folder, type='uv_unposed_no_shape', sizes=sizes)
+    downscale_uv_maps_unposed(folder=folder, type='uv_unposed_no_shape', sizes=sizes)
-    # downscale_uv_maps_curv(folder=folder, type='uv_curvature', sizes=sizes)
+    downscale_uv_maps_curv(folder=folder, type='uv_curvature', sizes=sizes)
-    # downscale_uv_maps_posed_body(folder=folder, type='uv_posed_body_no_orient', sizes=sizes)
+    downscale_uv_maps_posed_body(folder=folder, type='uv_posed_body_no_orient', sizes=sizes)
 
 
     print()

downscaling_test.py

@@ -125,7 +125,7 @@ if __name__ == '__main__':
 
 
     factory = DatasetFactory()
-    signature_dynamic = { 'id': (),
+    signature_dynamic = { 'id': (), 'frame_nr': (),
                           'pose_shape': (72 + 11,),
                         #   'uv_posed': (3, 512, 256,),
                         #   'uv_body_posed': (3, 512, 256,),

downscaling_test_curv.py

@@ -96,7 +96,7 @@ if __name__ == '__main__':
 
 
     factory = DatasetFactory()
-    signature_dynamic = { 'id': (),
+    signature_dynamic = { 'id': (), 'frame_nr': (),
                           'pose_shape': (72 + 11,),
                         #   'uv_posed': (3, 512, 256,),
                         #   'uv_body_posed': (3, 512, 256,),

edge_detection.py

@@ -0,0 +1,15 @@
+import os
+import cv2
+import numpy as np
+import plotly.express as ex
+
+if __name__ == '__main__':
+
+    img = np.load('/home/workspace/data/dataset/train/00042/uv_unposed_no_shape/Jumpsuit_58.npz')['uv_unposed_no_shape']
+    curv = np.load('/home/workspace/data/dataset/train/00042/curvature_uv/Jumpsuit_58.npz')['curvature_uv']
+    edges = cv2.Canny(image=(np.round((img+1) * 125)).astype(np.uint8), threshold1=25, threshold2=175, apertureSize=5, L2gradient=True)
+
+    ex.imshow(edges).show()
+    ex.imshow(curv).show()
+
+    print()