diff --git a/models/upscaler.py b/models/upscaler.py index 086df8b..100ebf9 100644 --- a/models/upscaler.py +++ b/models/upscaler.py @@ -99,11 +99,11 @@ class UpscalerBasicNew(tf.keras.Model): 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.Dense(4096, 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.Dense(2048, 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.Dense(2048, 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'), ]) @@ -123,6 +123,36 @@ class UpscalerBasicNew(tf.keras.Model): return x +class UpscalerBasicConv(tf.keras.Model): + + def __init__(self, low_res=32): + super(UpscalerBasicConv, self).__init__() + + self.low_res = low_res + + self.seq = tf.keras.Sequential([ + tf.keras.layers.Conv2D(4096, kernel_size=3, padding='same', activation='relu', name='upscale_enhancer_conv_1'), + tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_1'), + tf.keras.layers.Conv2D(2048, kernel_size=3, padding='same', activation='relu', name='upscale_enhancer_conv_2'), + tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_2'), + tf.keras.layers.Conv2D(1024, kernel_size=3, padding='same', activation='relu', name='upscale_enhancer_conv_3'), + tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_3'), + tf.keras.layers.Conv2D(128, kernel_size=3, padding='same', activation='relu', name='upscale_enhancer_conv_4'), + tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_4'), + tf.keras.layers.Dense(3, activation='tanh', name='upscale_enhancer_dense_final'), + ]) + + + def call(self, inputs, training=None): + img, mask = inputs + + x = self.seq(img, training=training) + + 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 ee30e6f..1818dbf 100644 --- a/upscale_test.py +++ b/upscale_test.py @@ -1,5 +1,5 @@ import os -os.environ["CUDA_VISIBLE_DEVICES"] = "1" +os.environ["CUDA_VISIBLE_DEVICES"] = "3" import tensorflow as tf gpus = tf.config.experimental.list_physical_devices('GPU') @@ -30,11 +30,14 @@ from dataset import DatasetFactory, StaticDataset from custom_utils.custom_utils import batch_to_dict # from custom_utils.smpl import SMPL # from models.upscaler import Upscaler -from models.upscaler import UpscalerBasicNew as Upscaler +from models.upscaler import UpscalerBasicConv as Upscaler -name='up_3_d256_512_2' +# !!!!!!!!!! CHANGE SSIM BASED ON RES +# !!!!!!!!!! CHANGE SKIRT OR NOT + +name='32_512_3' def init_wandb(): - run = wandb.init(project='vae', name=name, + run = wandb.init(project='upscale', name=name, notes='\ changed ms-ssim scale weights | \ inpainting style | \ @@ -44,6 +47,8 @@ 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 | \ +changed l1 loss reduce axis | \ +try convolutional architecture | \ ') artifact = wandb.Artifact('code', type='code') artifact.add_file('upscale_test.py') @@ -53,7 +58,6 @@ named the layers, added tanh to final layer | \ artifact.add_file('convnext/models/convnext_tf.py') run.log_artifact(artifact) -# !!!!!!!!!! CHANGE SKIRT OR NOT uv2mesh = UV2Mesh(res=256, is_skirt=False) @@ -98,17 +102,15 @@ 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' + low_resolution_modifier], batch['uv_mask' + low_resolution_modifier]]) + pred = model([batch['uv_unposed' + low_resolution_modifier], batch['uv_mask' + low_resolution_modifier]], training=True) 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' + 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)) + 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' + 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']) @@ -118,17 +120,17 @@ def train_step(batch): -# UPDATE! @tf.function def valid_step(batch): - pred = model(batch['uv_unposed' + low_resolution_modifier], batch['uv_mask']) + pred = model([batch['uv_unposed' + low_resolution_modifier], batch['uv_mask' + low_resolution_modifier]], training=False) + 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) )), axis=1)) - ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(tf.transpose(pred, [0,2,3,1]) + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=2)) - # ms_ssim = tf.reduce_mean(tf.image.ssim(tf.transpose(pred * uv_mask_body[tf.newaxis,tf.newaxis,:,:], [0,2,3,1]) + 1, tf.transpose(batch['uv_body_posed' + resolution_modifier] * uv_mask_body[tf.newaxis,tf.newaxis,:,:], [0,2,3,1]) + 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' + 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']) @@ -138,7 +140,7 @@ def valid_step(batch): def plot_everything(batch: dict, type: str): - pred = model([batch['uv_unposed' + low_resolution_modifier], batch['uv_mask' + low_resolution_modifier]]) + pred = model([batch['uv_unposed' + low_resolution_modifier], batch['uv_mask' + low_resolution_modifier]], training=False) pred = tf.image.resize(pred, size=(high_res, high_res//2), method=tf.image.ResizeMethod.LANCZOS5) mesh_pred = uv2mesh(pred) @@ -162,20 +164,20 @@ def plot_everything(batch: dict, type: str): def train(): - # valid_iterator = iter(dataset_valid) + valid_iterator = iter(dataset_valid) step = 0 for batch in dataset_train: batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic)) train_loss = train_step(batch) - # data_valid = valid_iterator.get_next() - # data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, signature_dynamic)) - valid_loss = {}#valid_step(data_valid) + data_valid = valid_iterator.get_next() + data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, signature_dynamic)) + valid_loss = valid_step(data_valid) - if step % 10000 == 0: + if step % 10000 == 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') wandb.log({ **train_loss, **valid_loss, **train_figs, **valid_figs }, step=step) else: wandb.log({ **train_loss, **valid_loss }, step=step) @@ -187,8 +189,8 @@ def train(): print(step) - if step > 300000: - break + # if step > 900000: + # break step += 1 @@ -203,7 +205,7 @@ if __name__ == '__main__': # 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 = 256 + low_res = 32 high_res = 512 uv_mask_body = tf.ones((low_res, low_res//2)) low_resolution_modifier = '_' + str(low_res) @@ -222,10 +224,10 @@ if __name__ == '__main__': 'uv_mask' + low_resolution_modifier: (low_res, low_res//2,), } - 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) + dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train', batch_size=8, signature=signature_dynamic, use_filled=False, resolutions=[high_res, low_res], in_memory=False) + static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train', signature=signature_static, use_filled=False, resolutions=[high_res, low_res], in_memory=True) + dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/valid', batch_size=8, signature=signature_dynamic, use_filled=False, resolutions=[high_res, low_res], in_memory=False) + static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/valid', signature=signature_static, use_filled=False, resolutions=[high_res, low_res], in_memory=True) model = Upscaler(low_res=low_res)