From 3c798b473d00fc783545279d9e2601462c247969 Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Tue, 4 Oct 2022 21:50:58 +0000 Subject: [PATCH] backup --- configs/config_pyramid.py | 77 ++++++++++++++++------- custom_utils/custom_utils.py | 4 +- cvae_pyramid.py | 114 +++++++++++++++++++++++------------ dataset.py | 8 ++- models/pyramid_level.py | 24 +++----- models/pyramid_model.py | 42 +++++++++++-- models/upscaler.py | 46 ++++++++++++-- upscale_test.py | 80 +++++++++++++----------- 8 files changed, 270 insertions(+), 125 deletions(-) diff --git a/configs/config_pyramid.py b/configs/config_pyramid.py index 513368d..10ea014 100644 --- a/configs/config_pyramid.py +++ b/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 diff --git a/custom_utils/custom_utils.py b/custom_utils/custom_utils.py index 82136bc..a94c9b5 100644 --- a/custom_utils/custom_utils.py +++ b/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: diff --git a/cvae_pyramid.py b/cvae_pyramid.py index ea32530..5dfa1fd 100644 --- a/cvae_pyramid.py +++ b/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: diff --git a/dataset.py b/dataset.py index b75aca6..274dc3d 100644 --- a/dataset.py +++ b/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) diff --git a/models/pyramid_level.py b/models/pyramid_level.py index c73a8a6..bc76f65 100644 --- a/models/pyramid_level.py +++ b/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 diff --git a/models/pyramid_model.py b/models/pyramid_model.py index dc3d056..bfa95a9 100644 --- a/models/pyramid_model.py +++ b/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 diff --git a/models/upscaler.py b/models/upscaler.py index c123174..086df8b 100644 --- a/models/upscaler.py +++ b/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): diff --git a/upscale_test.py b/upscale_test.py index dab34ea..ee30e6f 100644 --- a/upscale_test.py +++ b/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() \ No newline at end of file