diff --git a/upscale_test.py b/upscale_test.py index 1818dbf..93395d5 100644 --- a/upscale_test.py +++ b/upscale_test.py @@ -1,5 +1,5 @@ import os -os.environ["CUDA_VISIBLE_DEVICES"] = "3" +os.environ["CUDA_VISIBLE_DEVICES"] = "4" import tensorflow as tf gpus = tf.config.experimental.list_physical_devices('GPU') @@ -20,7 +20,7 @@ import wandb import random import cv2 import numpy as np -from custom_utils.custom_utils import UV2Mesh, load_model, maskFacesBatch +from custom_utils.custom_utils import UV2Mesh, get_summary_html, load_model, maskFacesBatch import plotly.graph_objects as go import plotly.express as ex @@ -35,7 +35,7 @@ from models.upscaler import UpscalerBasicConv as Upscaler # !!!!!!!!!! CHANGE SSIM BASED ON RES # !!!!!!!!!! CHANGE SKIRT OR NOT -name='32_512_3' +name='128_512_skirt' def init_wandb(): run = wandb.init(project='upscale', name=name, notes='\ @@ -49,7 +49,8 @@ training from low res to low res | \ named the layers, added tanh to final layer | \ changed l1 loss reduce axis | \ try convolutional architecture | \ -') +smaller features | \ +', dir='/home/workspace/local-data/temp') artifact = wandb.Artifact('code', type='code') artifact.add_file('upscale_test.py') artifact.add_file('models/upscaler.py') @@ -58,7 +59,12 @@ try convolutional architecture | \ artifact.add_file('convnext/models/convnext_tf.py') run.log_artifact(artifact) -uv2mesh = UV2Mesh(res=256, is_skirt=False) +is_skirt = True +is_cape = False +skirt_modifier = '_skirt' if is_skirt else '' +smpl_res = 3638 if is_cape else 14475 + +uv2mesh = UV2Mesh(res=256, is_skirt=is_skirt, is_smpl_orig=is_cape) @@ -109,8 +115,8 @@ def train_step(batch): 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' + 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)) + 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']) @@ -129,8 +135,8 @@ def valid_step(batch): 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' + 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)) + 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']) @@ -146,7 +152,7 @@ def plot_everything(batch: dict, type: str): mesh_pred = uv2mesh(pred) mesh_orig = batch['mesh_unposed'] - F = np.load("custom_utils/smpl_faces.npy") + F = np.load("custom_utils/smpl_faces" + skirt_modifier + ("_orig" if is_cape else "") + ".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])] @@ -175,22 +181,25 @@ def train(): data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, signature_dynamic)) valid_loss = valid_step(data_valid) - if step % 10000 == 0 or (step < 5000 and step % 500 == 0): + to_log = { **train_loss, **valid_loss, **{ 'current_step': step} } + if step % 5000 == 0 or (step < 5000 and step % 500 == 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) + to_log = { **to_log, **train_figs, **valid_figs } + + if step == 0: + to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } } - if step % 25000 == 0: + if step % 5000 == 0: manager.save(checkpoint_number=step) + wandb.log(to_log, step = step) print(step) - # if step > 900000: - # break + if step > 100000: + break step += 1 @@ -205,7 +214,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 = 32 + low_res = 128 high_res = 512 uv_mask_body = tf.ones((low_res, low_res//2)) low_resolution_modifier = '_' + str(low_res) @@ -216,26 +225,33 @@ if __name__ == '__main__': 'pose_shape': (72 + 11,), '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,), + 'mesh_unposed': (smpl_res, 3,), + 'mesh_posed': (smpl_res, 3,), } signature_static = { - 'mesh_mask': (14475,), + 'mesh_mask': (smpl_res,), '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', 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) + dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_skirt', batch_size=8, signature=signature_dynamic, use_filled=False, resolutions=[high_res, low_res], in_memory=False, is_cape=is_cape) + static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_skirt', signature=signature_static, use_filled=False, resolutions=[high_res, low_res], in_memory=True, is_cape=is_cape) + dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/valid_skirt', batch_size=8, signature=signature_dynamic, use_filled=False, resolutions=[high_res, low_res], in_memory=False, is_cape=is_cape) + static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/valid_skirt', signature=signature_static, use_filled=False, resolutions=[high_res, low_res], in_memory=True, is_cape=is_cape) - model = Upscaler(low_res=low_res) + # dataset_train = factory.get_dataset(folder='/home/workspace/local-data/cape_processed/train/female', batch_size=8, signature=signature_dynamic, use_filled=False, resolutions=[high_res, low_res], in_memory=False, is_cape=is_cape) + # static_dataset_train = StaticDataset(folder='/home/workspace/local-data/cape_processed/train/female', signature=signature_static, use_filled=False, resolutions=[high_res, low_res], in_memory=True, is_cape=is_cape) + # dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/cape_processed/test/female', batch_size=8, signature=signature_dynamic, use_filled=False, resolutions=[high_res, low_res], in_memory=False, is_cape=is_cape) + # static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/cape_processed/test/female', signature=signature_static, use_filled=False, resolutions=[high_res, low_res], in_memory=True, is_cape=is_cape) + + + model = Upscaler(low_res=low_res, kernel_size=5) 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) + manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + name, max_to_keep=20) # load_model(model, 'checkpoints/up_2_3') train() \ No newline at end of file