import os os.environ["CUDA_VISIBLE_DEVICES"] = "4" import tensorflow as tf gpus = tf.config.experimental.list_physical_devices('GPU') if gpus: try: # Currently, memory growth needs to be the same across GPUs for gpu in gpus: tf.config.experimental.set_memory_growth(gpu, True) logical_gpus = tf.config.experimental.list_logical_devices('GPU') print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs") except RuntimeError as e: # Memory growth must be set before GPUs have been initialized print(e) import time import argparse import wandb import random import cv2 import numpy as np from custom_utils.custom_utils import UV2Mesh, load_model, maskFacesBatch 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 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 UpscalerBasic as Upscaler name='up_2_3_skirt_d32' 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\ ') artifact = wandb.Artifact('code', type='code') artifact.add_file('upscale_test.py') artifact.add_file('models/upscaler.py') artifact.add_file('losses.py') artifact.add_file('dataset.py') artifact.add_file('convnext/models/convnext_tf.py') run.log_artifact(artifact) uv2mesh = UV2Mesh(res=256, is_skirt=True) def plot_3d(mesh_gt, mesh_pred_unposed, faces): mesh_pred_unposed[:,2] += 1 fig = go.Figure(data=[ go.Mesh3d( # 8 vertices of a cube x=mesh_gt[:,0], y=mesh_gt[:,2], z=mesh_gt[:,1], # i, j and k give the vertices of triangles i = faces[:,0], j = faces[:,1], k = faces[:,2] ), go.Mesh3d( # 8 vertices of a cube x=mesh_pred_unposed[:,0], y=mesh_pred_unposed[:,2], z=mesh_pred_unposed[:,1], # i, j and k give the vertices of triangles i = faces[:,0], j = faces[:,1], k = faces[:,2], ) ]) fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=2, z=1)) fig.update_layout( scene = dict( xaxis = dict(nticks=4, range=[-1,1],), yaxis = dict(nticks=4, range=[-1,3],), zaxis = dict(nticks=4, range=[-1,1],),) ) return fig @tf.function def train_step(batch): with tf.GradientTape() as tape: pred = model([batch['uv_unposed' + resolution_modifier], batch['uv_mask' + resolution_modifier]]) 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)) 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']) gradients = tape.gradient(loss, model.trainable_variables) optimizer.apply_gradients(zip(gradients, model.trainable_variables)) return { 'train/loss': loss, 'train/l1_loss': l1_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim } # UPDATE! @tf.function def valid_step(batch): pred = model(batch['uv_unposed' + resolution_modifier], batch['uv_mask']) 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)) 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']) return { 'valid/loss': loss, 'valid/l1_loss': l1_loss, 'valid/euclid_error': metric, 'valid/ms_ssim': ms_ssim } def plot_everything(batch: dict, type: str): pred = model([batch['uv_unposed' + resolution_modifier], batch['uv_mask' + resolution_modifier]]) mesh_pred = uv2mesh(pred) mesh_orig = batch['mesh_unposed'] F = np.load("custom_utils/smpl_faces_skirt.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])] gt = [mesh_orig[i].numpy()[mask[i], :] for i in range(pred.shape[0])] 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'] img = tf.concat((orig_img, pred), axis=2) imgs = [wandb.Image((img[i] + 1) / 2.0) for i in range(4)] return { **figs, **{ type + '/images': imgs} } def train(): # 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) if step % 2500 == 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: manager.save(checkpoint_number=step) print(step) step += 1 if __name__ == '__main__': init_wandb() allow_memory_growth() parser = argparse.ArgumentParser(description='') parser.add_argument('--name', type=str) args = vars(parser.parse_args()) 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 uv_mask_body = tf.ones((low_res, low_res//2)) resolution_modifier = '_' + str(low_res) factory = DatasetFactory() signature_dynamic = { 'id': (), 'pose_shape': (72 + 11,), 'uv_unposed' + resolution_modifier: (3, low_res, low_res//2,), 'uv_unposed': (3, 512, 256,), 'mesh_unposed': (14475, 3,), } signature_static = { 'mesh_mask': (14475,), 'uv_mask': (512, 256,), 'uv_mask' + 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_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) 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') train()