import os os.environ["CUDA_VISIBLE_DEVICES"] = "0" import time import argparse import wandb import random import cv2 import numpy as np import tensorflow as tf from custom_utils.custom_utils import UV2Mesh, 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.ae_body_model import AutoEncoder name='ae_body_8' def init_wandb(): run = wandb.init(project='vae', name=name, notes='convnext + convnext, add dense with tanh at the end of reverse convnext, larger features. add lr scheduler. fft3d instead of rfft3d. \ fft2d per channel and separate for top and bottom. fft_loss, all data. added convnext block as bottleneck, same as before, added extra metric. \ . lowered learning rate because it goes to nan. shallow third stage. slower lr decay. with validation. higher lr rate.\ two more downscaling than up, so decoder output is resized wto final resolution\ larger batch_size.\ flattened bottleneck with 128 latent_dim\ remove global orientation\ fixed gender issue with deoriented bodies, again, aaand again\ trained on 100 garments\ remove multiplications with garment uv mask') artifact = wandb.Artifact('code', type='code') artifact.add_file('auto_encoder_body.py') artifact.add_file('losses.py') artifact.add_file('models/ae_body_model.py') artifact.add_file('dataset.py') artifact.add_file('convnext/models/convnext_tf.py') run.log_artifact(artifact) uv2mesh = UV2Mesh(res=256) def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, collision_color): mesh_pred_unposed[:,1] += 1 collision[:,1] += 2 fig = go.Figure(data=[ go.Mesh3d( # 8 vertices of a cube x=mesh_gt[:,0], y=mesh_gt[:,1], z=mesh_gt[:,2], # 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[:,1], z=mesh_pred_unposed[:,2], # 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=collision[:,0], y=collision[:,1], z=collision[:,2], # i, j and k give the vertices of triangles i = collision_faces[:,0], j = collision_faces[:,1], k = collision_faces[:,2], vertexcolor=collision_color ) ]) 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_body_posed']) # pred = (pred + zero_body) * batch['mesh_mask'][:,:,tf.newaxis] mesh_pred = uv2mesh(pred) # mesh_pred = pred mesh_orig = batch['mesh_body_posed'] e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, tf.ones_like(batch['mesh_mask']))) c_loss = 0 # 1e-2 * tf.reduce_mean(collision_penalty(mesh_pred, batch['mesh_body_posed'], batch['mesh_mask'])) # l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(batch['mesh_unposed'] - pred)), axis=-1) * batch['mesh_mask']) l1_loss = tf.reduce_mean(tf.reduce_sum(((tf.math.abs(batch['uv_body_posed'] - pred) * uv_mask_body[tf.newaxis,tf.newaxis,:,:])), axis=1)) ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(tf.transpose(pred * uv_mask_body[tf.newaxis,tf.newaxis,:,:], [0,2,3,1]) + 1, tf.transpose(batch['uv_body_posed'] * uv_mask_body[tf.newaxis,tf.newaxis,:,:], [0,2,3,1]) + 1, max_val=2)) fft_loss = 0 # tf.reduce_mean(tf.reduce_sum(tf.math.abs(tf.signal.fft3d(tf.complex(pred, tf.zeros_like(pred))) - tf.signal.fft3d(tf.complex(batch['uv_body_posed'], tf.zeros_like(batch['uv_body_posed'])))), axis=[1,2,3])) # loss = fft_loss_pc # + e_loss + c_loss loss = 1e3 * (1 - ms_ssim) mesh_orig_2 = uv2mesh(batch['uv_body_posed']) # mesh_skinned = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred) 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']) metric_2 = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig_2 - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask']) euclid_posed = 0 # tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_skinned - batch['mesh_posed']) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask']) gradients = tape.gradient(loss, model.trainable_variables) # gradients = [tf.clip_by_norm(g, 1.0) for g in gradients] # gradients = [tf.clip_by_value(g, -1.0, 1.0) for g in gradients] optimizer.apply_gradients(zip(gradients, model.trainable_variables)) return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/collision_loss': c_loss, 'train/l1_loss': l1_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim, 'train/fft_loss': fft_loss, 'train/euclid_error_reconst': metric_2, 'train/euclid_error_posed': euclid_posed } # UPDATE! @tf.function def valid_step(batch): pred = model([batch['uv_body_posed'], batch['uv_mask']]) mesh_pred = uv2mesh(pred) mesh_orig = batch['mesh_body_posed'] e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask'])) c_loss = 0 #1e-2 * tf.reduce_mean(collision_penalty(mesh_pred, batch['mesh_body_posed'], batch['mesh_mask'])) l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(batch['uv_body_posed'] - pred)), axis=1) * batch['uv_mask']) ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(tf.transpose(pred, [0,2,3,1]) + 1, tf.transpose(batch['uv_body_posed'], [0,2,3,1]) + 1, max_val=2)) fft_loss = 0 #tf.reduce_mean(tf.reduce_sum(tf.math.abs(tf.signal.fft3d(tf.complex(pred, tf.zeros_like(pred))) - tf.signal.fft3d(tf.complex(batch['uv_body_posed'], tf.zeros_like(batch['uv_body_posed'])))), axis=[1,2,3])) loss = 1e3 * (1 - ms_ssim) mesh_orig_2 = uv2mesh(batch['uv_body_posed']) 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']) metric_2 = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig_2 - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask']) return { 'valid/loss': loss, 'valid/edge_loss': e_loss, 'valid/collision_loss': c_loss, 'valid/l1_loss': l1_loss, 'valid/euclid_error': metric, 'valid/ms_ssim': ms_ssim, 'valid/fft_loss': fft_loss, 'valid/euclid_error_reconst': metric_2 } def plot_everything(batch: dict, type: str): pred = model(batch['uv_body_posed']) mesh_pred = uv2mesh(pred) mesh_orig = batch['mesh_body_posed'] orig_img = tf.transpose(batch['uv_body_posed'], [0,2,3,1]) pred_img = tf.transpose(pred, [0,2,3,1]) img = tf.concat((orig_img, pred_img), axis=2) imgs = [wandb.Image((img[i] + 1) / 2.0) for i in range(4)] F = np.load("custom_utils/smpl_faces.npy") mask = batch['mesh_mask'].numpy().astype(bool) # garment_unposed, faces = maskFacesBatch(mesh_pred.numpy(), F, mask) garment_unposed, faces = mesh_pred.numpy(), F # gt = [mesh_orig[i].numpy()[mask[i], :] for i in range(pred.shape[0])] gt = [mesh_orig[i].numpy()[:, :] for i in range(pred.shape[0])] body = batch['mesh_body_posed'].numpy() # mesh_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred).numpy() # garment_posed = [mesh_posed[i, mask[i], :] for i in range(body.shape[0])] garment_posed = garment_unposed collision = [np.concatenate((body[i], garment_posed[i]), axis=0) for i in range(len(garment_posed))] # collision_faces = [np.concatenate((F, faces[i] + 14475), axis=0) for i in range(len(garment_posed))] collision_faces = [np.concatenate((F, faces + 14475), axis=0) for i in range(len(garment_posed))] collision_color = [np.concatenate((np.ones_like(body[i], dtype=np.int32)*100, np.ones_like(garment_posed[i], dtype=np.int32)*230), axis=0) for i in range(len(garment_posed))] figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], faces, collision[i], collision_faces[i], collision_color[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) } 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) factory = DatasetFactory() signature_dynamic = { 'id': (), 'pose_shape': (72 + 11,), # 'uv_posed': (3, 512, 256,), 'uv_body_posed': (3, 512, 256,), 'mesh_body_posed': (14475, 3,), # 'mesh_posed': (14475, 3,), # 'mesh_unposed': (14475, 3,), # 'uv_unposed': (3, 512, 256,) } signature_static = { #'uv_static': (3, 512, 256,), # 'mesh_static': (14475, 3,), 'mesh_mask': (14475,), # 'uv_mask': (512, 256,), # 'uv_shaped_body': (3, 512, 256,), } dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100', batch_size=32, signature=signature_dynamic, use_filled=True) static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', signature=signature_static, use_filled=True) dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/valid', batch_size=32, signature=signature_dynamic) static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/valid', signature=signature_static) model = AutoEncoder() # lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(0.0005, decay_steps=30000, decay_rate=0.1, staircase=False) 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) train()