import os import time import argparse from cv2 import polarToCart import wandb import random import numpy as np import tensorflow as tf from custom_utils.custom_utils import UV2Mesh, maskFacesBatch from tensorboardX import SummaryWriter import plotly.graph_objects as go from utils import str_to_bool from tf_utils import check_tf_version, allow_memory_growth, split_gpu_for_testing from load_models import load_generator, load_discriminator from losses import d_logistic, d_logistic_r1_reg, g_logistic_non_saturating, g_logistic_ns_pathreg, edge_loss, uv_mask_loss # from losses import collision_penalty as collision_penalty from losses import collision_loss as collision_penalty from convnext.models.convnext_tf import ConvNeXt from stylegan2_ae.generator import Generator os.environ["CUDA_VISIBLE_DEVICES"] = "2" run = wandb.init(config=tf.compat.v1.flags.FLAGS, sync_tensorboard=True, project='stylegan', name='ae_unposed_1', notes='convnext + stylegan generator') artifact = wandb.Artifact('code', type='code') artifact.add_file('auto_encoder.py') artifact.add_file('stylegan2_ae/layers/mapping_block.py') run.log_artifact(artifact) uv2mesh = UV2Mesh(res=256) class DatasetSimple(tf.data.Dataset): def _generator(folder): # Opening the file folder = folder.decode("utf-8") uv_folder = os.path.join(folder, 'uv_posed') frames = os.listdir(uv_folder) random.shuffle(frames) for frame in frames: id, frame_nr = frame.split('.')[0].split('_') mesh_mask = np.load(os.path.join(folder, 'mesh_mask', id + '.npy')); mesh_mask = mesh_mask.astype(np.float32) uv_mask = np.load(os.path.join(folder, 'uv_mask', id + '.npz')); uv_mask=uv_mask['uv_mask']; uv_mask = uv_mask.astype(np.float32) mesh_unposed = np.load(os.path.join(folder, 'mesh_unposed', id + '_' + frame_nr + '.npz')); mesh_unposed = mesh_unposed['mesh_unposed'] uv_unposed = np.load(os.path.join(folder, 'uv_unposed', id + '_' + frame_nr + '.npz')); uv_unposed = uv_unposed['uv_unposed']; uv_unposed = np.transpose(uv_unposed, [2, 0, 1]) yield uv_unposed, mesh_unposed, uv_mask, mesh_mask def __new__(cls, folder, res=256): return tf.data.Dataset.from_generator( cls._generator, output_signature = (tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32), tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32), tf.TensorSpec(shape = (res * 2, res,), dtype = tf.float32), tf.TensorSpec(shape = (14475,), dtype = tf.float32), ), args=(folder, ) ) def filter_resolutions_featuremaps(resolutions, featuremaps, res): index = resolutions.index(res) filtered_resolutions = resolutions[:index + 1] filtered_featuremaps = featuremaps[:index + 1] return filtered_resolutions, filtered_featuremaps def plot_3d(mesh_gt, mesh_pred, faces): mesh_pred[:,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[:,0], y=mesh_pred[:,1], z=mesh_pred[:,2], # 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 class AutoEncoder(tf.keras.Model): def __init__(self): super(AutoEncoder, self).__init__() resolutions = [4, 8, 16, 32, 64, 128, 256, 512, 1024] g_featuremaps = [256, 256, 256, 256, 256, 128, 64, 32, 16] train_resolutions, train_g_featuremaps = filter_resolutions_featuremaps(resolutions, g_featuremaps, 512) g_params = { 'z_dim': 256, 'w_dim': 256, 'labels_dim': -1, 'n_mapping': 8, 'resolutions': train_resolutions, 'featuremaps': train_g_featuremaps, } # test_latent = tf.ones((1, g_params['z_dim']), dtype=tf.float32) self.decoder = Generator(g_params) # _ = generator(test_latent) self.encoder = ConvNeXt(depths=[3, 3, 27, 3], dims=[96, 192, 256, g_params['z_dim']], include_top=False) self.pooling = tf.keras.layers.GlobalAveragePooling2D() def call(self, inputs, training=None): x, uv_mask = inputs x = tf.transpose(x, [0,2,3,1]) x = self.encoder(x, training=training) x = self.pooling(x) x = self.decoder(x, training=training) x = x * uv_mask[:,tf.newaxis,:,:] return x @tf.function def train_step(batch): uv_unposed, mesh_unposed, uv_mask, mesh_mask = batch with tf.GradientTape() as tape: pred = model([uv_unposed, uv_mask]) mesh_pred = uv2mesh(pred) mesh_orig = mesh_unposed e_loss = 0.0#edge_loss(mesh_pred, mesh_orig, mesh_mask) c_loss = 0.0#collision_penalty(mesh_pred, mesh_body_posed, mesh_mask) l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(uv_unposed - pred)), axis=1) * uv_mask) loss = l1_loss + e_loss + c_loss metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(mesh_mask[:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(mesh_mask) gradients = tape.gradient(loss, model.trainable_variables) optimizer.apply_gradients(zip(gradients, model.trainable_variables)) return { 'loss': loss, 'edge_loss': e_loss, 'collision_loss': c_loss, 'l1_loss': l1_loss, 'euclid_error': metric } @tf.function def valid_step(batch): uv_unposed, mesh_unposed, uv_mask, mesh_mask = batch pred = model([uv_unposed, uv_mask]) mesh_pred = uv2mesh(pred) mesh_orig = mesh_unposed e_loss = 0.0#edge_loss(mesh_pred, mesh_orig, mesh_mask) c_loss = 0.0#collision_penalty(mesh_pred, mesh_body_posed, mesh_mask) l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(uv_unposed - pred)), axis=1) * uv_mask) loss = l1_loss + e_loss + c_loss metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(mesh_mask[:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(mesh_mask) return { 'loss': loss, 'edge_loss': e_loss, 'collision_loss': c_loss, 'l1_loss': l1_loss, 'euclid_error': metric } if __name__ == '__main__': allow_memory_growth() name = 'ae_unposed_1' dataset_train = DatasetSimple('/home/workspace/data/dataset/stylegan_new/train').cache().batch(8).repeat() dataset_valid = DatasetSimple('/home/workspace/data/dataset/stylegan_new/valid').cache().batch(8).repeat() model = AutoEncoder() # optimizer = tf.optimizers.Adam(learning_rate=0.002, beta_1=0.0, beta_2=0.99, epsilon=1e-8) optimizer = tf.optimizers.Adam() folder = 'models/ae_logs/' train_summary_writer = tf.summary.create_file_writer(folder + name + '/train') # valid_summary_writer = tf.summary.create_file_writer(folder + name + '/valid') step = 0 for batch in dataset_train: train_loss = train_step(batch) for key in train_loss.keys(): with train_summary_writer.as_default(): tf.summary.scalar(key, train_loss[key], step=step) data_valid = batch # data_valid = iter(dataset_valid).get_next() # valid_loss = valid_step(data_valid) # for key in valid_loss.keys(): # with valid_summary_writer.as_default(): # tf.summary.scalar(key, valid_loss[key], step=step) if step % 250 == 0: pred = model([data_valid[0], data_valid[2]]) mesh_pred = uv2mesh(pred) mesh_orig = data_valid[1] orig_img = tf.transpose(data_valid[0], [0,2,3,1]) pred_img = tf.transpose(pred, [0,2,3,1]) img = tf.concat((orig_img, pred_img), axis=2) with train_summary_writer.as_default(): tf.summary.image('images', img, step=step) F = np.load("custom_utils/smpl_faces.npy") garment, faces = maskFacesBatch(mesh_pred.numpy(), F, data_valid[3].numpy().astype(bool)) gt = [mesh_orig[i].numpy()[data_valid[3][i].numpy().astype(bool), :] for i in range(pred.shape[0])] for i in range(pred.shape[0]): fig = plot_3d(gt[i], garment[i], faces[i]) wandb.log({"garment_" + str(i): fig}) print(step) step += 1