backup

configs/config_pyramid.py

@@ -1,27 +1,38 @@
 from configs.config_base import config
 
-config['gpu'] = '1'
-config['name'] = 'cvae_pyr_4_4'
+config['gpu'] = '2'
+config['name'] = 'cvae_pyr_6_32_64_128_3_x'
 config['notes'] = '\
-offsets over average unposed garment. remove global orientation (in posed body input and in pose_shape parameters. \
-retrained body encoder with deoriented bodies \
-use mean l2 loss. small batch. l2 * 1000, balancing to 1 from 0.01\
-fixed gender issue and retrained body encoder, FINAL\
-remove shape from unposed garment\
-training on 100 garments with offsets over the average of these\
-fixed autoencoders!!!!!!!!!\
-common condition encoders for all levels, only run once\
-full res loss on each level\
-all higher levels are added to the output of first level\
-using custom MLP upscale enhancement\
-layers 32 and 512\
-try loading weights from previous trainings\
+offsets over average unposed garment. remove global orientation (in posed body input and in pose_shape parameters. | \
+retrained body encoder with deoriented bodies | \
+use mean l2 loss. small batch. l2 * 1000, balancing to 1 from 0.01 | \
+fixed gender issue and retrained body encoder, FINAL | \
+remove shape from unposed garment | \
+training on 100 garments with offsets over the average of these | \
+fixed autoencoders!!!!!!!!! | \
+common condition encoders for all levels, only run once | \
+full res loss on each level | \
+all higher levels are added to the output of first level | \
+using custom MLP upscale enhancement | \
+layers 32, 64, 128 | \
+testing plotting | \
+only train highest level, all other use from pretrained pyramids with one less level | \
+maybe upscale enhancer was not frozen? --> fixed pretrained approach | \
+upscale from level to level, not to final res, only at the end to final res  | \
+updated upscale enhancers | \
+try an old trained model as pretraining for new level | \
+ \
 '
-config['dataset_in_memory'] = False
+
+# DON'T FORGET TO RESET SSIM LOSS USAGE
+# DON'T FORGET TO CHECK MODEL LOADING
+
+config['dataset_in_memory'] = True
 config['resolution'] = None
 
 config['lr'] = 0.0001
 config['starting_balancing_coeff'] = 0.0001
+config['reg_loss_coeff'] = 0 # set to 0 if don't want to use edge+collision
 
 config['signature_dynamic'] = { 'id': (),
                           'pose_shape': (72 + 11,),
@@ -41,7 +52,9 @@ config['pyramid'] = dict()
 config['pyramid']['is_label'] = False
 config['pyramid']['latent_dim'] = config['latent_dim']
 
-config['pyramid']['resolutions'] = [32, 512]
+config['pyramid']['resolutions'] = [32, 64, 128]
+config['pyramid']['loss_weights'] = [1,1,1]
+config['pyramid']['pretrained_pyramid'] = 'checkpoints/cvae_pyr_6_32_64'
 
 for res in config['pyramid']['resolutions']:
   if res != 512:
@@ -53,12 +66,16 @@ for res in config['pyramid']['resolutions']:
       config['signature_static']['uv_mask' + res_modifier] = (res, res//2)
 
 
-config['pyramid']['static_model'] = 'checkpoints/ae_static_7_2_x'
-config['pyramid']['body_model'] = 'checkpoints/ae_body_9_x'
-config['pyramid']['curvature_model'] = 'checkpoints/ae_curv_9_x'
+# config['pyramid']['static_model'] = 'checkpoints/ae_static_7_2_x'
+# config['pyramid']['body_model'] = 'checkpoints/ae_body_9_x'
+# config['pyramid']['curvature_model'] = 'checkpoints/ae_curv_9_x'
+config['pyramid']['static_model'] = ''
+config['pyramid']['body_model'] = ''
+config['pyramid']['curvature_model'] = ''
 config['pyramid']['freeze_encoders'] = True
 
 config['avg_garment_path'] = ''
+config['pyramid']['final_upscale_enhancer'] = 'checkpoints/up_3_d128_256'
 
 
 config['pyramid']['levels'] = []
@@ -87,7 +104,22 @@ for res in config['pyramid']['resolutions']:
     level['decoder_depths'] = [16, 16, 9, 9]
     level['decoder_dims'] = [256, 256, 128, 64]
     level['avg_garment_path'] = ''
-    level['upscale_enhancer'] = 'checkpoints/up_2_3_d128'
+    level['upscale_enhancer'] = ''
+    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'] = [27, 27, 16]
+    level['encoder_dims'] = [256, 256, 256]
+    level['decoder_input'] = (4, 2, 256)
+    level['decoder_depths'] = [27, 27, 16]
+    level['decoder_dims'] = [256, 256, 256]
+    level['avg_garment_path'] = ''
+    level['upscale_enhancer'] = 'checkpoints/up_3_d64_128'
+    level['upscale_enhancer_freeze'] = False
     config['pyramid']['levels'].append(level)
   elif res == 32:
     level = dict()
@@ -100,7 +132,8 @@ for res in config['pyramid']['resolutions']:
     level['decoder_depths'] = [27, 27]
     level['decoder_dims'] = [256, 256]
     level['avg_garment_path'] = ''
-    level['upscale_enhancer'] = 'checkpoints/up_2_3'
+    level['upscale_enhancer'] = 'checkpoints/up_3_d32_64'
+    level['upscale_enhancer_freeze'] = True
     config['pyramid']['levels'].append(level)
   else:
     raise NotImplementedError

custom_utils/custom_utils.py

@@ -214,8 +214,8 @@ def load_model(model, checkpoint, disable_assert=False, assert_consumed_on=False
     ckpt = tf.train.Checkpoint(model=model)
     manager = tf.train.CheckpointManager(ckpt, checkpoint, max_to_keep=1)
     if not disable_assert:
-        assert manager.latest_checkpoint is not None
-        assert manager.latest_checkpoint.split('/')[-1] != 'ckpt-0'
+        assert manager.latest_checkpoint is not None, checkpoint
+        assert manager.latest_checkpoint.split('/')[-1] != 'ckpt-0', checkpoint
     if assert_consumed_on:
         ckpt.restore(manager.latest_checkpoint).assert_consumed()
     else:

cvae_pyramid.py

@@ -27,8 +27,13 @@ import plotly.express as ex
 from plotly.subplots import make_subplots
 from PIL import Image
 
+if config['is_skirt']:
+    from losses import edge_loss_skirt as edge_loss
+else:
+    from losses import edge_loss
+
 from tf_utils import check_tf_version, allow_memory_growth, split_gpu_for_testing
-from losses import edge_loss, uv_mask_loss, bend_loss, kl_loss
+from losses import 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
@@ -48,14 +53,14 @@ def init_wandb():
     artifact.add_file('models/upscaler.py')
     artifact.add_file('losses.py')
     artifact.add_file('configs/config_base.py')
-    artifact.add_file('configs/config_cvae.py')
+    artifact.add_file('configs/config_pyramid.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)
+uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'])
 
 def plot_3d(mesh_gt, mesh_pred_unposed_levels, faces, mesh_posed, collision, collision_faces, collision_color):
     fig = make_subplots(rows=1, cols=1+len(mesh_pred_unposed_levels)+1+1, start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(1+len(mesh_pred_unposed_levels)+1+1)] ])
@@ -148,15 +153,16 @@ def group_inputs_for_levels(batch):
     for res in config['pyramid']['resolutions']:
         res_modifier = '_' + str(res) if res != 512 else ''
         levels_input.append([batch['uv_unposed' + res_modifier], batch['uv_static' + res_modifier], batch['uv_body_posed' + res_modifier], batch['uv_curvature' + res_modifier], batch['uv_mask' + res_modifier], batch['uv_mask']])
-    return levels_input
+    conds = [batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature']]
+    return levels_input, conds
 
 
 @tf.function
 def train_step(batch):
 
     with tf.GradientTape() as tape:
-        levels_input = group_inputs_for_levels(batch)
-        out_levels = model(levels_input, training=True)
+        levels_input, conds = group_inputs_for_levels(batch)
+        out_levels = model([levels_input, conds], training=True)
 
         shape = batch['pose_shape'][:, 72:]
         gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
@@ -166,32 +172,37 @@ def train_step(batch):
         output = dict()
         reconstr_loss_total = 0
         kl_total = 0
-        for i, res in enumerate(config['pyramid']['resolutions']):
-            res_modifier = '_' + str(res) if res != 512 else ''
-            pred, mean, logvar = out_levels[i]
 
-            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, res in enumerate(config['pyramid']['resolutions']):
+        i = -1
+        res = config['pyramid']['resolutions'][-1]
+        res_modifier = '_' + str(res) if res != 512 else ''
+        pred, mean, logvar = out_levels[i]
 
-            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']))
-            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 + 1, batch['uv_unposed'] + 1, max_val=3, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
-            kl = kl_loss(mean, logvar)
+        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)
 
-            reconstr_loss_total += l2_loss #+ e_loss + c_loss
-            kl_total += kl
+        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']))
+        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 - tf.reduce_min(pred), batch['uv_unposed'] - tf.reduce_min(pred), max_val=tf.reduce_max(pred) - tf.reduce_min(pred), 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))
+        kl = kl_loss(mean, logvar)
 
-            mean = tf.reduce_mean(mean)
-            logvar = tf.reduce_mean(logvar)
+        reconstr_loss_total += config['pyramid']['loss_weights'][i] * (l2_loss + (e_loss + c_loss) * config['reg_loss_coeff'])
+        # reconstr_loss_total += config['pyramid']['loss_weights'][i] * (l2_loss + (e_loss + c_loss) * config['reg_loss_coeff'])
+        kl_total += kl
 
-            metric = 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'])
+        mean = tf.reduce_mean(mean)
+        logvar = tf.reduce_mean(logvar)
 
-            output = { **output, **{ 'train/edge_loss' + res_modifier: e_loss, 'train/collision_loss' + res_modifier: c_loss,'train/euclid_error' + res_modifier: metric, 
-                                     'train/ms_ssim' + res_modifier: ms_ssim, '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/euclid_error': metric } }
+        metric = 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'])
+
+        output = { **output, **{ 'train/edge_loss' + res_modifier: e_loss, 'train/collision_loss' + res_modifier: c_loss,'train/euclid_error' + res_modifier: metric, 
+                                    'train/ms_ssim' + res_modifier: ms_ssim, '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/euclid_error': metric } }
 
         loss = reconstr_loss_total + balancing_coeff * kl_total
 
@@ -248,19 +259,20 @@ def plot_everything(batch: dict, type: str):
     mesh_orig_posed = batch['mesh_posed'].numpy()
     body = batch['mesh_body_posed'].numpy()
     
-    orig_img = batch['uv_unposed']
-    uv_static_img = batch['uv_static']
-    uv_body_posed_img = batch['uv_body_posed']
-    uv_curvature = np.tile(batch['uv_curvature'].numpy().argmax(-1)[:,:,:,np.newaxis], [1,1,1,3]) / 2 - 1
-    pred_img = np.concatenate([output_levels[i][0] for i in range(len(output_levels))], axis=2)
+    orig_img = (batch['uv_unposed'].numpy() + 1) / 2
+    uv_static_img = (batch['uv_static'].numpy() + 1) / 2
+    uv_body_posed_img = (batch['uv_body_posed'].numpy() + 1) / 2
+    uv_curvature = (np.tile(batch['uv_curvature'].numpy().argmax(-1)[:,:,:,np.newaxis], [1,1,1,3]) / 6)
+    pred_img = np.concatenate([tf.image.resize(output_levels[i][0], (512, 256), tf.image.ResizeMethod.LANCZOS5) for i in range(len(output_levels))], axis=2)
+    pred_img = (pred_img - pred_img.min(axis=(1,2,3))[:,np.newaxis,np.newaxis,np.newaxis]) / (pred_img.max(axis=(1,2,3))[:,np.newaxis,np.newaxis,np.newaxis] - pred_img.min(axis=(1,2,3))[:,np.newaxis,np.newaxis,np.newaxis])
     img = tf.concat((uv_static_img, uv_body_posed_img, uv_curvature, orig_img, pred_img), axis=2)
-    imgs = [wandb.Image((img[i] + 1) / 2.0) for i in range(4)]
+    imgs = [wandb.Image(img[i]) 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])]).numpy()
 
-    F = np.load("custom_utils/smpl_faces.npy")
+    F = np.load("custom_utils/smpl_faces" + skirt_modifier + ".npy")
     mask = batch['mesh_mask'].numpy().astype(bool)
 
     _, faces = maskFacesBatch(np.ones_like(mesh_orig), F, mask)
@@ -419,11 +431,15 @@ if __name__ == '__main__':
     args = vars(parser.parse_args())
 
     smpl = SMPL()
+    skirt_modifier = '_skirt' if config['is_skirt'] else ''
+    dataset_resolutions = config['pyramid']['resolutions'].copy()
+    if 512 not in dataset_resolutions:
+        dataset_resolutions.append(512)
 
     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['pyramid']['resolutions'])
-    static_dataset_train = StaticDataset(folder=config['train_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=config['pyramid']['resolutions'])
-    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['pyramid']['resolutions'])
+    dataset_train = factory.get_dataset(folder=config['train_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=dataset_resolutions, in_memory=config['dataset_in_memory'])
+    static_dataset_train = StaticDataset(folder=config['train_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=dataset_resolutions)
+    dataset_valid = factory.get_dataset(folder=config['valid_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=dataset_resolutions, in_memory=config['dataset_in_memory'])
     static_dataset_valid = StaticDataset(folder=config['valid_dir'], signature=config['signature_static'])
 
     model: Pyramid = Pyramid(config)
@@ -440,9 +456,29 @@ if __name__ == '__main__':
     ckpt = tf.train.Checkpoint(model=model)
     manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + config['name'], max_to_keep=3)
 
-    load_model(model.cond_encoder, 'checkpoints/cvae_19')
-    load_model(model.levels[0], 'checkpoints/cvae_low_4_3_d32')
-    load_model(model.levels[-1], 'checkpoints/cvae_19')
+    load_model(model, config['pyramid']['pretrained_pyramid'])
+    model.static_enc.trainable = False
+    model.body_enc.trainable = False
+    model.curv_enc.trainable = False
+    model.cond_encoder.trainable = False
+
+    # the last level doesn't have an upscale enhancer as it's enhanced for the final resolutin directly
+    # might now be necessary, as enhancing to final or a bit higher resolution seems to be the same
+    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
+    load_model(model.upscale_enhancer, config['pyramid']['final_upscale_enhancer'])
+    model.upscale_enhancer.trainable = False
+
+    load_model(model.levels[-1], 'checkpoints/cvae_low_2_d128_cont')
+
+    for i in range(len(model.levels)-1):
+        model.levels[i].trainable = False
+
+    # load_model(model.cond_encoder, 'checkpoints/cvae_low_4_3_d32')
+    # load_model(model.levels[0], 'checkpoints/cvae_low_4_3_d32')
+    # load_model(model.levels[-1], 'checkpoints/cvae_19')
+
 
     step = 0
     for batch in dataset_train:

dataset.py

@@ -59,8 +59,10 @@ class StaticDataset():
         self.signature = signature
         self.folder = folder
         self.in_memory = in_memory
-        assert set(resolutions).issubset(set([32, 128, 512]))
-        self.possible_keys = ['uv_static', 'uv_static_32', 'uv_static_128', 'mesh_static', 'mesh_mask', 'uv_mask', 'uv_mask_32', 'uv_mask_128', 'uv_shaped_body']
+        assert set(resolutions).issubset(set([32, 64, 128, 256, 512]))
+        self.possible_keys = ['uv_static', 'uv_static_32', 'uv_static_64', 'uv_static_128', 'uv_static_256',
+                              'uv_mask', 'uv_mask_32', 'uv_mask_64', 'uv_mask_128', 'uv_mask_256',
+                              'mesh_static', 'mesh_mask', 'uv_shaped_body']
         self.keys = [value for value in signature.keys() if value in self.possible_keys]
         if 'uv_static_128' in signature.keys() and 'uv_static_128' not in self.keys:
             self.keys.append('uv_static_128')
@@ -146,7 +148,7 @@ class DatasetFactory():
     
     def get_dataset(self, folder, batch_size, signature, shuffle=True, use_filled=False, resolutions=[512], in_memory=True):
         # assert (use_filled and len(resolutions) >= 1) or (not use_filled and len(resolutions) == 1)
-        assert set(resolutions).issubset(set([32, 128, 512]))
+        assert set(resolutions).issubset(set([32, 64, 128, 256, 512]))
 
         # static_data = StaticDataset(folder, signature)
         dataset_idx = DatasetIdx(folder, shuffle)

models/pyramid_level.py

@@ -9,7 +9,7 @@ import scipy as sp
 
 from custom_utils.custom_utils import load_model
 from models.base import AEEncoder, VariationalEncoder, Decoder
-from models.upscaler import UpscalerBasic
+from models.upscaler import UpscalerBasic, UpscalerBasicNew
 
 
 class PyramidLevel(tf.keras.Model):
@@ -24,13 +24,9 @@ class PyramidLevel(tf.keras.Model):
 
         self.upscale_enhancer = None
         if config['upscale_enhancer'] != '':
-            self.upscale_enhancer = UpscalerBasic(512, config['resolution'])
+            self.upscale_enhancer = UpscalerBasicNew(config['resolution'])
             load_model(self.upscale_enhancer, config['upscale_enhancer'])
-            self.upscale_enhancer = tf.keras.Sequential([
-                self.upscale_enhancer.flatten,
-                self.upscale_enhancer.seq,
-                self.upscale_enhancer.reshape
-                ])
+            self.upscale_enhancer.trainable = not config['upscale_enhancer_freeze']
 
         if config['avg_garment_path'] != '':
             self.avg_garment = np.load(config['avg_garment_path'])
@@ -39,10 +35,8 @@ class PyramidLevel(tf.keras.Model):
 
 
     def enhance(self, low, mask, training=None):
-        x = self.upscale_enhancer(low, training=training)
-        x = x * (1 - mask[:,:,:,tf.newaxis])
-        low = x + low
-        return low
+        x = self.upscale_enhancer([low, mask], training=training)
+        return x
 
 
     def sample(self, z_cond, uv_mask, noise=None):
@@ -87,11 +81,11 @@ class PyramidLevel(tf.keras.Model):
 
         x = x * uv_mask[:,:,:,tf.newaxis]
 
-        if self.upscale_enhancer is not None:
-            pred_low = self.enhance(x, uv_mask, training=training)
-            x = tf.image.resize(pred_low, (512, 256), tf.image.ResizeMethod.LANCZOS5)
+        # if self.upscale_enhancer is not None:
+        #     pred_low = self.enhance(x, uv_mask, training=training)
+        #     x = tf.image.resize(pred_low, (512, 256), tf.image.ResizeMethod.LANCZOS5)
         
-        x = x * uv_mask_high[:,:,:,tf.newaxis]
+        # x = x * uv_mask_high[:,:,:,tf.newaxis]
 
         return x, mean, logvar
 

models/pyramid_model.py

@@ -10,6 +10,7 @@ from custom_utils.custom_utils import load_model
 from models.base import AEEncoder
 from models.cvae_model import CVAE
 from models.pyramid_level import PyramidLevel
+from models.upscaler import UpscalerBasic, UpscalerBasicNew
 
 
 class Pyramid(tf.keras.Model):
@@ -41,6 +42,16 @@ class Pyramid(tf.keras.Model):
 
         self.levels = [PyramidLevel({ **config['pyramid']['levels'][i], **{ 'cond_ae': config['cond_ae']} }) for i in range(len(config['pyramid']['levels']))]
 
+ 
+        self.upscale_enhancer = None
+        if config['pyramid']['final_upscale_enhancer'] != '':
+            self.upscale_enhancer = UpscalerBasicNew(low_res=config['pyramid']['resolutions'][-1])
+            load_model(self.upscale_enhancer, config['pyramid']['final_upscale_enhancer'])
+            self.upscale_enhancer.trainable = False
+
+    def enhance(self, low, mask, training=None):
+        x = self.upscale_enhancer([low, mask], training=training)
+        return x
 
     def encode_conds(self, uv_template, uv_body_posed, uv_curv):
         z_static = self.static_enc(uv_template)
@@ -53,18 +64,37 @@ class Pyramid(tf.keras.Model):
         return z_cond
 
     def call(self, inputs, training=None):
-        uv_static, uv_body, uv_curv = inputs[-1][1], inputs[-1][2], inputs[-1][3]
+        levels_input, conds = inputs
+        uv_static, uv_body, uv_curv = conds[0], conds[1], conds[2]
         z_cond = self.encode_conds(uv_static, uv_body, uv_curv)
 
         out_levels = []
-        pred, mean, logvar = self.levels[0](inputs[0], z_cond, training=training)
-        out_levels.append([pred, mean, logvar])
+        pred, mean, logvar = self.levels[0](levels_input[0], z_cond, training=training)
+        pred = self.levels[0].enhance(pred, levels_input[0][-2], training=training)  # levels_input[i][-2]=uv_mask_res
+        if 1 != len(self.levels):
+            pred = tf.image.resize(pred, (self.levels[1].size, self.levels[1].size//2), tf.image.ResizeMethod.LANCZOS5)
+            out_levels.append([pred, mean, logvar])
+            pred = pred * levels_input[1][-2][:,:,:,tf.newaxis]
+        else:
+            pred = tf.image.resize(pred, (512, 256), tf.image.ResizeMethod.LANCZOS5)
+            out_levels.append([pred, mean, logvar])
 
         for i in range(1, len(self.levels)):
-            pred_level, mean, logvar = self.levels[i](inputs[i], z_cond, training=training)
+            pred_level, mean, logvar = self.levels[i](levels_input[i], z_cond, training=training)
             pred = pred + pred_level
-            out_levels.append([pred, mean, logvar])
-        
+            if i != len(self.levels) - 1:
+                pred = self.levels[i].enhance(pred, levels_input[i][-2], training=training)
+                out_levels.append([pred, mean, logvar])
+                pred = tf.image.resize(pred, (self.levels[i+1].size, self.levels[i+1].size//2), tf.image.ResizeMethod.LANCZOS5)
+                pred = pred * levels_input[i+1][-2][:,:,:,tf.newaxis]
+            else:
+                if self.upscale_enhancer is not None:
+                    pred = self.enhance(pred, levels_input[-1][-2], training=training)
+                    pred = tf.image.resize(pred, (512, 256), tf.image.ResizeMethod.LANCZOS5)
+                out_levels.append([pred, mean, logvar])
+                pred = pred * levels_input[-1][-1][:,:,:,tf.newaxis] # full res uv mask
+
+
         return out_levels
 
 

models/upscaler.py

@@ -55,9 +55,11 @@ class Upscaler(tf.keras.Model):
 
 class UpscalerBasic(tf.keras.Model):
     
-    def __init__(self, bottleneck_size, size=32):
+    def __init__(self, bottleneck_size, low_res=32, high_res=64):
         super(UpscalerBasic, self).__init__()
 
+        self.low_res = low_res
+        self.high_res = high_res
         self.flatten = tf.keras.layers.Flatten()
 
         self.seq = tf.keras.Sequential([
@@ -67,10 +69,10 @@ class UpscalerBasic(tf.keras.Model):
                     tf.keras.layers.BatchNormalization(),
                     tf.keras.layers.Dense(1024, activation='relu'),
                     tf.keras.layers.BatchNormalization(),
-                    tf.keras.layers.Dense(size * (size//2) * 3),
+                    tf.keras.layers.Dense(low_res * (low_res//2) * 3),
         ])
 
-        self.reshape = tf.keras.layers.Reshape((size, (size//2), 3))
+        self.reshape = tf.keras.layers.Reshape((low_res, (low_res//2), 3))
 
     def call(self, inputs, training=None):
         img, mask = inputs
@@ -81,10 +83,46 @@ class UpscalerBasic(tf.keras.Model):
 
         x = x * (1 - mask[:,:,:,tf.newaxis])
         x = x + img
-        x = tf.image.resize(x, size=(512, 256), method=tf.image.ResizeMethod.LANCZOS5)
+        x = tf.image.resize(x, size=(self.high_res, self.high_res//2), method=tf.image.ResizeMethod.LANCZOS5)
 
         return x
 
+
+
+
+class UpscalerBasicNew(tf.keras.Model):
+    
+    def __init__(self, low_res=32):
+        super(UpscalerBasicNew, self).__init__()
+
+        self.low_res = low_res
+        self.flatten = tf.keras.layers.Flatten(name='upscale_enhancer_flattent')
+
+        self.seq = tf.keras.Sequential([
+                    tf.keras.layers.Dense(1024, activation='relu', name='upscale_enhancer_dense_1'),
+                    tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_1'),
+                    tf.keras.layers.Dense(1024, activation='relu', name='upscale_enhancer_dense_2'),
+                    tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_2'),
+                    tf.keras.layers.Dense(1024, activation='relu', name='upscale_enhancer_dense_3'),
+                    tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_3'),
+                    tf.keras.layers.Dense(low_res * (low_res//2) * 3, activation='tanh', name='upscale_enhancer_dense_final'),
+        ])
+
+        self.reshape = tf.keras.layers.Reshape((low_res, (low_res//2), 3), name='upscale_enhancer_reshape')
+
+    def call(self, inputs, training=None):
+        img, mask = inputs
+
+        x = self.flatten(img)
+        x = self.seq(x, training=training)
+        x = self.reshape(x)
+
+        x = x * (1 - mask[:,:,:,tf.newaxis])
+        x = x + img
+
+        return x
+
+
 # class Upscaler(tf.keras.Model):
     
 #     def __init__(self, bottleneck_size):

upscale_test.py

@@ -1,5 +1,5 @@
 import os
-os.environ["CUDA_VISIBLE_DEVICES"] = "4"
+os.environ["CUDA_VISIBLE_DEVICES"] = "1"
 import tensorflow as tf
 
 gpus = tf.config.experimental.list_physical_devices('GPU')
@@ -25,24 +25,26 @@ import plotly.graph_objects as go
 import plotly.express as ex
 
 from tf_utils import check_tf_version, allow_memory_growth
-from losses import edge_loss, uv_mask_loss
+# from losses import edge_loss, uv_mask_loss
 from dataset import DatasetFactory, StaticDataset
 from custom_utils.custom_utils import batch_to_dict
-from custom_utils.smpl import SMPL
+# from custom_utils.smpl import SMPL
 # from models.upscaler import Upscaler
-from models.upscaler import UpscalerBasic as Upscaler
+from models.upscaler import UpscalerBasicNew as Upscaler
 
-name='up_2_3_skirt_d32'
+name='up_3_d256_512_2'
 def init_wandb():
     run = wandb.init(project='vae', name=name,
             notes='\
-changed ms-ssim scale weights\
-inpainting style\
-MLP based\
-smaller batch\
-use mask for loss\
-started from up_2_3\
-    ')
+changed ms-ssim scale weights | \
+inpainting style | \
+MLP based | \
+smaller batch | \
+use mask for loss | \
+using ms ssim for higher res | \
+training from low res to low res | \
+named the layers, added tanh to final layer | \
+')
     artifact = wandb.Artifact('code', type='code')
     artifact.add_file('upscale_test.py')
     artifact.add_file('models/upscaler.py')
@@ -51,7 +53,8 @@ started from up_2_3\
     artifact.add_file('convnext/models/convnext_tf.py')
     run.log_artifact(artifact)
 
-uv2mesh = UV2Mesh(res=256, is_skirt=True)
+# !!!!!!!!!! CHANGE SKIRT OR NOT
+uv2mesh = UV2Mesh(res=256, is_skirt=False)
 
 
 
@@ -95,14 +98,17 @@ def plot_3d(mesh_gt, mesh_pred_unposed, faces):
 @tf.function
 def train_step(batch):
     with tf.GradientTape() as tape:
-        pred = model([batch['uv_unposed' + resolution_modifier], batch['uv_mask' + resolution_modifier]])
+        pred = model([batch['uv_unposed' + low_resolution_modifier], batch['uv_mask' + low_resolution_modifier]])
+        pred = tf.image.resize(pred, size=(high_res, high_res//2), method=tf.image.ResizeMethod.LANCZOS5)
 
         mesh_pred = uv2mesh(pred)
         mesh_orig = batch['mesh_unposed']
 
-        l1_loss = tf.reduce_mean(tf.reduce_sum(((tf.math.abs(batch['uv_unposed'] - pred * batch['uv_mask'][:,:,:,tf.newaxis]) )), axis=1))
-        # ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred * batch['uv_mask'][:,:,:,tf.newaxis] + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=2, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
-        ms_ssim = tf.reduce_mean(tf.image.ssim(pred * batch['uv_mask'][:,:,:,tf.newaxis] + 1, batch['uv_unposed'] * batch['uv_mask'][:,:,:,tf.newaxis] + 1, max_val=2))
+        l1_loss = tf.reduce_mean(tf.reduce_sum(((tf.math.abs(batch['uv_unposed' + high_resolution_modifier] - pred * batch['uv_mask' + high_resolution_modifier][:,:,:,tf.newaxis]) )), axis=1))
+        # ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred * batch['uv_mask'][:,:,:,tf.newaxis] + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=2, power_factors=(0.5, 0.2, 0.2, 0.1)))
+        # ms_ssim = tf.reduce_mean(tf.image.ssim(pred * batch['uv_mask'][:,:,:,tf.newaxis] + 1, batch['uv_unposed'] * batch['uv_mask'][:,:,:,tf.newaxis] + 1, max_val=2))
+        ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred * batch['uv_mask' + high_resolution_modifier][:,:,:,tf.newaxis] + 1, batch['uv_unposed' + high_resolution_modifier] * batch['uv_mask' + high_resolution_modifier][:,:,:,tf.newaxis] + 1, max_val=2, power_factors=(0.5, 0.2, 0.2, 0.1)))
+        # ms_ssim = tf.reduce_mean(tf.image.ssim(pred * batch['uv_mask' + high_resolution_modifier][:,:,:,tf.newaxis] + 1, batch['uv_unposed' + high_resolution_modifier] * batch['uv_mask' + high_resolution_modifier][:,:,:,tf.newaxis] + 1, max_val=2))
         loss = 1e3 * (1 - ms_ssim)
 
     metric = 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'])
@@ -115,7 +121,7 @@ def train_step(batch):
 # UPDATE!
 @tf.function
 def valid_step(batch):
-    pred = model(batch['uv_unposed' + resolution_modifier], batch['uv_mask'])
+    pred = model(batch['uv_unposed' + low_resolution_modifier], batch['uv_mask'])
 
     mesh_pred = uv2mesh(pred)
     mesh_orig = batch['mesh_unposed']
@@ -132,12 +138,13 @@ def valid_step(batch):
 
 def plot_everything(batch: dict, type: str):
 
-    pred = model([batch['uv_unposed' + resolution_modifier], batch['uv_mask' + resolution_modifier]])
+    pred = model([batch['uv_unposed' + low_resolution_modifier], batch['uv_mask' + low_resolution_modifier]])
+    pred = tf.image.resize(pred, size=(high_res, high_res//2), method=tf.image.ResizeMethod.LANCZOS5)
     
     mesh_pred = uv2mesh(pred)
     mesh_orig = batch['mesh_unposed']
     
-    F = np.load("custom_utils/smpl_faces_skirt.npy")
+    F = np.load("custom_utils/smpl_faces.npy")
     mask = batch['mesh_mask'].numpy().astype(bool)
     garment_unposed, faces = maskFacesBatch(mesh_pred.numpy(), F, mask)
     # gt = [mesh_orig[i].numpy()[mask[i], :] for i in range(pred.shape[0])]
@@ -146,7 +153,7 @@ def plot_everything(batch: dict, type: str):
     figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], faces[i]) for i in range(4) }
     # 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) }
     
-    orig_img = batch['uv_unposed']
+    orig_img = batch['uv_unposed' + high_resolution_modifier]
     img = tf.concat((orig_img, pred), axis=2)
     imgs = [wandb.Image((img[i] + 1) / 2.0) for i in range(4)]
 
@@ -166,20 +173,23 @@ def train():
         # data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, signature_dynamic))
         valid_loss = {}#valid_step(data_valid)
 
-        if step % 2500 == 0:
+        if step % 10000 == 0:
             train_figs = plot_everything(batch, 'train')
             valid_figs = {}#plot_everything(data_valid, 'valid')
             wandb.log({ **train_loss, **valid_loss, **train_figs, **valid_figs }, step=step)
         else:
             wandb.log({ **train_loss, **valid_loss }, step=step)
             
-        if step % 10000 == 0:
+        if step % 25000 == 0:
             manager.save(checkpoint_number=step)
 
 
 
         print(step)
 
+        if step > 300000:
+            break
+
         step += 1
 
 if __name__ == '__main__':
@@ -189,39 +199,41 @@ if __name__ == '__main__':
     parser.add_argument('--name', type=str)
     args = vars(parser.parse_args())
     
-    smpl = SMPL()
+    # smpl = SMPL()
     # 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)
 
-    low_res = 32
+    low_res = 256
+    high_res = 512
     uv_mask_body = tf.ones((low_res, low_res//2))
-    resolution_modifier = '_' + str(low_res)
+    low_resolution_modifier = '_' + str(low_res)
+    high_resolution_modifier = ('_' + str(high_res)) if high_res != 512 else ''
 
     factory = DatasetFactory()
     signature_dynamic = { 'id': (),
                           'pose_shape': (72 + 11,),
-                          'uv_unposed' + resolution_modifier: (3, low_res, low_res//2,),
-                          'uv_unposed': (3, 512, 256,),
+                          'uv_unposed' + low_resolution_modifier: (low_res, low_res//2, 3),
+                          'uv_unposed' + high_resolution_modifier: (high_res, high_res//2, 3),
                           'mesh_unposed': (14475, 3,),
                         }
     signature_static = { 
                          'mesh_mask': (14475,),
-                         'uv_mask': (512, 256,),
-                         'uv_mask' + resolution_modifier: (low_res, low_res//2,),
+                         'uv_mask' + high_resolution_modifier: (high_res, high_res//2,),
+                         'uv_mask' + low_resolution_modifier: (low_res, low_res//2,),
                         }
 
-    dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_skirt_25_2', batch_size=8, signature=signature_dynamic, use_filled=False, resolutions=[512, low_res])
-    static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_skirt_25_2', signature=signature_static, use_filled=False, resolutions=[512, low_res])
+    dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100', batch_size=8, signature=signature_dynamic, use_filled=False, resolutions=[high_res, low_res])
+    static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', signature=signature_static, use_filled=False, resolutions=[high_res, low_res])
     dataset_valid = None #factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/valid', batch_size=32, signature=signature_dynamic)
     static_dataset_valid = None #StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/valid', signature=signature_static)
 
-    model = Upscaler(bottleneck_size=512, size=low_res)
+    model = Upscaler(low_res=low_res)
 
     optimizer = tf.optimizers.Adam(learning_rate=0.0001)
     folder = 'models/ae_logs/'
 
     ckpt = tf.train.Checkpoint(model=model)
     manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + name, max_to_keep=3)
-    load_model(model, 'checkpoints/up_2_3')
+    # load_model(model, 'checkpoints/up_2_3')
 
     train()