From 94e693d7f63713b9654a7f5f93ded9bc26f7d624 Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Fri, 10 Jun 2022 15:32:25 +0000 Subject: [PATCH] plotting not working, losses fixed? --- .gitignore | 4 +- custom_utils/custom_utils.py | 24 +++++++++- inference_custom.py | 27 +++++++---- losses.py | 17 ++++--- train.py | 90 ++++++++++++++++++++++-------------- 5 files changed, 108 insertions(+), 54 deletions(-) diff --git a/.gitignore b/.gitignore index ee1e6da..4d09f44 100644 --- a/.gitignore +++ b/.gitignore @@ -134,4 +134,6 @@ dmypy.json # vscode .vscode/ -wandb/ \ No newline at end of file +wandb/ +output/ +output.zip \ No newline at end of file diff --git a/custom_utils/custom_utils.py b/custom_utils/custom_utils.py index b611790..aa151e8 100644 --- a/custom_utils/custom_utils.py +++ b/custom_utils/custom_utils.py @@ -62,4 +62,26 @@ class UV2Mesh(): def __call__(self, uv): return self.uv_to_mesh(uv) - \ No newline at end of file + + + +def _maskFaces(V, F, mask): + V_mask = mask + F_mask = np.take(V_mask, F.reshape(-1), axis=0) + F_mask = F_mask.reshape((-1, 3)).sum(axis=1) == 3 + + V_ind = np.zeros(V_mask.shape) + V_ind[V_mask] = np.arange(V_mask.sum()) + F_new = np.take(V_ind, F[F_mask, :].reshape(-1), axis=0).reshape((-1, 3)) + V_new = V[V_mask, :] + + return V_new, F_new + +def maskFacesBatch(V, F, mask): + Vs = [] + Fs = [] + for i in range(mask.shape[0]): + v, f = _maskFaces(V[i], F, mask[i]) + Vs.append(v) + Fs.append(f) + return Vs, Fs \ No newline at end of file diff --git a/inference_custom.py b/inference_custom.py index 9e48c15..08788b2 100644 --- a/inference_custom.py +++ b/inference_custom.py @@ -5,6 +5,7 @@ import sys import cv2 import numpy as np from PIL import Image +from dataset import DatasetSimple from load_models import load_generator from tf_utils import allow_memory_growth @@ -38,7 +39,7 @@ def uv_to_mesh(uv): return mesh -def inference(ckpt_dir, use_custom_cuda, res, folder=None, samples=1000): +def inference(ckpt_dir, use_custom_cuda, res, folder, data_folder, samples=100): # create generator resolutions = [4, 8, 16, 32, 64, 128, 256, 512, 1024] featuremaps = [512, 512, 512, 512, 512, 256, 128, 64, 32] @@ -52,17 +53,23 @@ def inference(ckpt_dir, use_custom_cuda, res, folder=None, samples=1000): 'featuremaps': featuremaps[:filter_index + 1], } generator = load_generator(g_params, is_g_clone=True, ckpt_dir=ckpt_dir, custom_cuda=use_custom_cuda) + dataset = DatasetSimple('/home/workspace/data/dataset/stylegan_new/train').cache().batch(1).repeat() - for i in range(samples): + i = 0 + for batch in dataset: print(i) # generate image fake_images = generator([tf.random.normal(shape=[1, g_params['z_dim']]), - tf.random.normal(shape=[1, g_params['labels_dim']])], + batch[1]], training=False, truncation_psi=0.5) - np.savez_compressed(os.path.join(folder, str(i+100) + '.npz'), uv=fake_images) + np.savez_compressed(os.path.join(folder, str(i) + '.npz'), uv=tf.squeeze(fake_images), template=tf.squeeze(batch[1][2]), body=tf.squeeze(batch[3]), mask=tf.squeeze(batch[4])) img = np.transpose(np.round((fake_images + 1) * 128)[0], [1,2,0]) - cv2.imwrite(os.path.join(folder, str(i+100) + '.jpg'), img) + cv2.imwrite(os.path.join(folder, str(i) + '.jpg'), img) + i += 1 + + if i > samples: + break return @@ -73,9 +80,10 @@ def main(): checkpoints = [ { 'res': 512, - 'ckpt_dir': '/workspace/Checkpoints/checkpoint_52000/stylegan2-test-2-512x512', - 'use_custom_cuda': False, - 'folder': '/workspace/Checkpoints/checkpoint_52000/output', + 'ckpt_dir': 'models/stylegan2-test-5-512x512', + 'use_custom_cuda': True, + 'folder': 'output', + 'data_folder': '/home/workspace/data/dataset/stylegan_new/train' }, ] @@ -85,9 +93,10 @@ def main(): use_custom_cuda = run_item['use_custom_cuda'] folder = run_item['folder'] message = f'{res}x{res} with custom cuda' if use_custom_cuda else f'{res}x{res} without custom cuda' + data_folder = run_item['data_folder'] print(message) - inference(ckpt_dir, use_custom_cuda, res, folder) + inference(ckpt_dir, use_custom_cuda, res, folder, data_folder) return diff --git a/losses.py b/losses.py index 017f8f8..c924b7a 100644 --- a/losses.py +++ b/losses.py @@ -21,7 +21,7 @@ def d_logistic(real_images, generator, discriminator, z_dim, conditions): d_loss += tf.math.softplus(-real_scores) - return fake_images, d_loss + return d_loss def d_logistic_r1_reg(real_images, generator, discriminator, z_dim, conditions): @@ -46,7 +46,7 @@ def d_logistic_r1_reg(real_images, generator, discriminator, z_dim, conditions): real_grads = r1_tape.gradient(real_loss, real_images) r1_penalty = tf.reduce_sum(tf.math.square(real_grads), axis=[1, 2, 3]) r1_penalty = tf.expand_dims(r1_penalty, axis=1) - return fake_images, d_loss, r1_penalty + return d_loss, r1_penalty def g_logistic_non_saturating(real_images, generator, discriminator, z_dim, conditions): @@ -104,7 +104,7 @@ F = np.load("custom_utils/smpl_faces.npy") def _nearest_neighbour(V, B): tree = cKDTree(B.numpy()) - return tree.query(V.numpy(), n_jobs=-1)[1] + return tree.query(V.numpy(), workers=-1)[1] @tf.function def _tf_nn_parallel(V, B): @@ -117,13 +117,14 @@ def collision_loss(V, B, mask, thr=.004): # compute correspondences idx = _tf_nn_parallel(V, B) N, M = V.shape[:2] - _idx = tf.tile(np.array(list(range(N)), np.int32)[:,None], [1, M]) + # _idx = tf.tile(np.array(list(range(N)), np.int32)[:,None], [1, M]) + _idx = tf.ones((N, M), dtype=tf.int32) * np.array(list(range(N)), np.int32)[:,None] idx = tf.stack((_idx, idx), axis=-1) # compute loss D = V - tf.gather_nd(B, idx) b_vn = tf.gather_nd(B_vn, idx) dot = tf.einsum('abc,abc->ab', D, b_vn) - loss = tf.reduce_sum(tf.minimum(dot * mask - thr, 0) ** 2) + loss = tf.reduce_sum((tf.minimum(dot - thr, 0) * mask) ** 2) return loss @@ -150,7 +151,6 @@ def edge_loss(x, y, mask): return loss - class PairwiseDistance(tf.keras.layers.Layer): def __init__(self, **kwargs): super(PairwiseDistance, self).__init__(**kwargs) @@ -191,7 +191,6 @@ def collision_penalty(va, vb, mask, eps=1e-3): if __name__ == '__main__': V = tf.ones((4, 14475, 3)) B = tf.zeros((4, 14475, 3)) - B_vn = vertex_normals(B, tf.tile(F[None], [B.shape[0], 1, 1])) - loss_1 = collision_loss(V, B, thr=.004, stop_gradient=False) - loss_2 = collision_penalty(V, B, B_vn) + m = tf.ones((4, 14475)) + loss_1 = collision_loss(V, B, m, thr=.004) print() diff --git a/train.py b/train.py index 98d71f6..1d56fc2 100644 --- a/train.py +++ b/train.py @@ -5,20 +5,20 @@ import wandb from cv2 import drawFrameAxes import numpy as np import tensorflow as tf -from custom_utils.custom_utils import UV2Mesh +from custom_utils.custom_utils import UV2Mesh, maskFacesBatch +from tensorboardX import SummaryWriter 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 dataset_ffhq import get_ffhq_dataset from dataset import DatasetSimple from losses import d_logistic, d_logistic_r1_reg, g_logistic_non_saturating, g_logistic_ns_pathreg, edge_loss # from losses import collision_penalty as collision_penalty from losses import collision_loss as collision_penalty -os.environ["CUDA_VISIBLE_DEVICES"] = "0" -wandb.init(config=tf.compat.v1.flags.FLAGS, sync_tensorboard=True, project='stylegan', name='test_conditioned_1') +os.environ["CUDA_VISIBLE_DEVICES"] = "1" +wandb.init(config=tf.compat.v1.flags.FLAGS, sync_tensorboard=True, project='stylegan', name='test_conditioned_2') def initiate_models(g_params, d_params, use_custom_cuda): @@ -120,54 +120,54 @@ class Trainer(object): return params_copy def d_train_step(self, dist_inputs): - real_images, conditions, mesh_static, mesh_body_posed, mask = dist_inputs[0] + real_images, conditions, _, _, _ = dist_inputs[0] with tf.GradientTape() as d_tape: # compute losses - fake_images, d_loss = d_logistic(real_images, self.generator, self.discriminator, self.g_params['z_dim'], conditions=conditions) + d_loss = d_logistic(real_images, self.generator, self.discriminator, self.g_params['z_dim'], conditions=conditions) - mesh_pred = uv2mesh(fake_images) + # mesh_pred = uv2mesh(fake_images) - e_loss = edge_loss(mesh_pred, mesh_static, mask) - c_loss = collision_penalty(mesh_pred, mesh_body_posed) + # e_loss = edge_loss(mesh_pred, mesh_static, mask) + # c_loss = collision_penalty(mesh_pred, mesh_body_posed) # scale loss d_loss = tf.reduce_sum(d_loss) * self.global_batch_scaler - e_loss = tf.reduce_sum(e_loss) * self.global_batch_scaler - c_loss = tf.reduce_sum(c_loss) * self.global_batch_scaler + # e_loss = tf.reduce_sum(e_loss) * self.global_batch_scaler + # c_loss = tf.reduce_sum(c_loss) * self.global_batch_scaler - d_loss = d_loss + e_loss + c_loss + d_loss = d_loss# + e_loss + c_loss d_gradients = d_tape.gradient(d_loss, self.discriminator.trainable_variables) self.d_optimizer.apply_gradients(zip(d_gradients, self.discriminator.trainable_variables)) - return d_loss, e_loss, c_loss + return d_loss#, e_loss, c_loss def d_train_step_reg(self, dist_inputs): - real_images, conditions, mesh_static, mesh_body_posed, mask = dist_inputs[0] + real_images, conditions, _, _, _ = dist_inputs[0] with tf.GradientTape() as d_tape: # compute losses - fake_images, d_gan_loss, r1_penalty = d_logistic_r1_reg(real_images, self.generator, self.discriminator, + d_gan_loss, r1_penalty = d_logistic_r1_reg(real_images, self.generator, self.discriminator, self.g_params['z_dim'], conditions=conditions) r1_penalty = r1_penalty * (0.5 * self.r1_gamma) * self.d_opt['reg_interval'] - mesh_pred = uv2mesh(fake_images) + # mesh_pred = uv2mesh(fake_images) - e_loss = edge_loss(mesh_pred, mesh_static, mask) - c_loss = collision_penalty(mesh_pred, mesh_body_posed) + # e_loss = edge_loss(mesh_pred, mesh_static, mask) + # c_loss = collision_penalty(mesh_pred, mesh_body_posed) # scale losses r1_penalty = tf.reduce_sum(r1_penalty) * self.global_batch_scaler d_gan_loss = tf.reduce_sum(d_gan_loss) * self.global_batch_scaler - e_loss = tf.reduce_sum(e_loss) * self.global_batch_scaler - c_loss = tf.reduce_sum(c_loss) * self.global_batch_scaler + # e_loss = tf.reduce_sum(e_loss) * self.global_batch_scaler + # c_loss = tf.reduce_sum(c_loss) * self.global_batch_scaler # combine - d_loss = d_gan_loss + r1_penalty + e_loss + c_loss + d_loss = d_gan_loss + r1_penalty# + e_loss + c_loss d_gradients = d_tape.gradient(d_loss, self.discriminator.trainable_variables) self.d_optimizer.apply_gradients(zip(d_gradients, self.discriminator.trainable_variables)) - return d_loss, d_gan_loss, r1_penalty, e_loss, c_loss + return d_loss, d_gan_loss, r1_penalty#, e_loss, c_loss def g_train_step(self, dist_inputs): real_images, conditions, mesh_static, mesh_body_posed, mask = dist_inputs[0] @@ -179,7 +179,7 @@ class Trainer(object): mesh_pred = uv2mesh(fake_images) e_loss = edge_loss(mesh_pred, mesh_static, mask) - c_loss = collision_penalty(mesh_pred, mesh_body_posed) + c_loss = collision_penalty(mesh_pred, mesh_body_posed, mask) # scale loss g_loss = tf.reduce_sum(g_loss) * self.global_batch_scaler @@ -205,7 +205,7 @@ class Trainer(object): mesh_pred = uv2mesh(fake_images) e_loss = edge_loss(mesh_pred, mesh_static, mask) - c_loss = collision_penalty(mesh_pred, mesh_body_posed) + c_loss = collision_penalty(mesh_pred, mesh_body_posed, mask) # scale loss pl_penalty = tf.reduce_sum(pl_penalty) * self.global_batch_scaler @@ -224,10 +224,10 @@ class Trainer(object): @tf.function def dist_d_train_step(inputs): per_replica_losses = strategy.run(fn=self.d_train_step, args=(inputs,)) - mean_d_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[0], axis=None) - mean_e_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[1], axis=None) - mean_c_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[2], axis=None) - return mean_d_loss, mean_e_loss, mean_c_loss + mean_d_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses, axis=None) + # mean_e_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[1], axis=None) + # mean_c_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[2], axis=None) + return mean_d_loss#, mean_e_loss, mean_c_loss @tf.function def dist_d_train_step_reg(inputs): @@ -235,9 +235,9 @@ class Trainer(object): mean_d_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[0], axis=None) mean_d_gan_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[1], axis=None) mean_r1_penalty = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[2], axis=None) - mean_e_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[3], axis=None) - mean_c_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[4], axis=None) - return mean_d_loss, mean_d_gan_loss, mean_r1_penalty, mean_e_loss, mean_c_loss + # mean_e_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[3], axis=None) + # mean_c_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[4], axis=None) + return mean_d_loss, mean_d_gan_loss, mean_r1_penalty#, mean_e_loss, mean_c_loss @tf.function def dist_g_train_step(inputs): @@ -270,6 +270,8 @@ class Trainer(object): # setup tensorboards train_summary_writer = tf.summary.create_file_writer(self.ckpt_dir) + mesh_writer = SummaryWriter(os.path.join(self.ckpt_dir)) + # loss metrics metric_d_loss = tf.keras.metrics.Mean('d_loss', dtype=tf.float32) @@ -290,9 +292,9 @@ class Trainer(object): # d train step if (step + 1) % self.d_opt['reg_interval'] == 0: - d_loss, d_gan_loss, r1_penalty, e_loss, c_loss = dist_d_train_step_reg((batch, )) + d_loss, d_gan_loss, r1_penalty = dist_d_train_step_reg((batch, )) else: - d_loss, e_loss, c_loss = dist_d_train_step((batch, )) + d_loss = dist_d_train_step((batch, )) d_gan_loss = d_loss r1_penalty = zero @@ -348,6 +350,9 @@ class Trainer(object): with train_summary_writer.as_default(): tf.summary.image('images', summary_image, step=step) + # mesh = uv2mesh(tf.transpose(summary_image[:, :self.train_res, :, :], [0,3,1,2])) + # plot_meshes(mesh_writer, garment=mesh, template=batch[1][2], body=batch[3], mask=batch[4], step=step) + # train_summary_writer.flush() # print every self.print_steps @@ -397,6 +402,23 @@ def filter_resolutions_featuremaps(resolutions, featuremaps, res): return filtered_resolutions, filtered_featuremaps +def plot_meshes(writer, garment, template, body, mask, step, nr_offset=0): + F = np.load("custom_utils/smpl_faces.npy") + garment, faces = maskFacesBatch(garment.numpy(), F, mask.astype(bool)) + template, _ = maskFacesBatch(template.numpy(), F, mask.astype(bool)) + body = body.numpy() + + collision = [np.concatenate((body[i], garment[i]), axis=0) for i in range(len(garment))] + collision_faces = [np.concatenate((F, faces + 14475), axis=0) for i in range(len(garment))] + + for i in range(len(garment)): + writer.add_mesh('garment_pc_pred_' + str(i+nr_offset), vertices=garment[i][np.newaxis,:,:], global_step=step) + writer.add_mesh('garment_pred_' + str(i+nr_offset), vertices=garment[i][np.newaxis,:,:], faces=faces[i][np.newaxis,:,:], global_step=step) + writer.add_mesh('template_' + str(i+nr_offset), vertices=template[i][np.newaxis,:,:], faces=faces[i][np.newaxis,:,:], global_step=step) + writer.add_mesh('body_' + str(i+nr_offset), vertices=body[i][np.newaxis,:,:], faces=F[np.newaxis,:,:], global_step=step) + writer.add_mesh('collision_' + str(i+nr_offset), vertices=collision[i][np.newaxis,:,:], faces=collision_faces[np.newaxis,:,:], global_step=step) + + def main(): # global program arguments parser parser = argparse.ArgumentParser(description='') @@ -482,7 +504,7 @@ def main(): 'train_res': args['train_res'], } - trainer = Trainer(training_parameters, name=f'stylegan2-test-4-{args["train_res"]}x{args["train_res"]}') + trainer = Trainer(training_parameters, name=f'stylegan2-test-5-{args["train_res"]}x{args["train_res"]}') trainer.train(dist_dataset, strategy) return