diff --git a/custom_utils/custom_utils.py b/custom_utils/custom_utils.py index 37d13a7..71a1e98 100644 --- a/custom_utils/custom_utils.py +++ b/custom_utils/custom_utils.py @@ -1,5 +1,6 @@ import numpy as np import tensorflow as tf +import plotly.graph_objects as go def m2uv_array(m2uv): @@ -85,4 +86,53 @@ def maskFacesBatch(V, F, mask): v, f = _maskFaces(V[i], F, mask[i]) Vs.append(v) Fs.append(f) - return Vs, Fs \ No newline at end of file + return Vs, Fs + + +def plot_3d(mesh_gt, mesh_pred, faces, collision, collision_faces, collision_color): + mesh_pred[:,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[:,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], + ), + 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 \ No newline at end of file diff --git a/requierements.txt b/requierements.txt index 51dd24e..8574fc9 100644 --- a/requierements.txt +++ b/requierements.txt @@ -4,4 +4,5 @@ opencv-python wandb OpenEXR tensorflow-graphics -scipy \ No newline at end of file +scipy +plotly \ No newline at end of file diff --git a/train.py b/train.py index 53d92c7..416ddfa 100644 --- a/train.py +++ b/train.py @@ -5,7 +5,7 @@ import wandb from cv2 import drawFrameAxes import numpy as np import tensorflow as tf -from custom_utils.custom_utils import UV2Mesh, maskFacesBatch +from custom_utils.custom_utils import UV2Mesh, maskFacesBatch, plot_3d from tensorboardX import SummaryWriter from utils import str_to_bool @@ -18,7 +18,7 @@ from losses import collision_loss as collision_penalty os.environ["CUDA_VISIBLE_DEVICES"] = "3" -run = wandb.init(config=tf.compat.v1.flags.FLAGS, sync_tensorboard=True, project='stylegan', name='test_conditioned_13(test-16)', notes='learning offsets, all losses from the start') +run = wandb.init(config=tf.compat.v1.flags.FLAGS, sync_tensorboard=True, project='stylegan', name='test_conditioned_19', notes='learning offsets, all losses from the start, smaller generator (from 128), slower learning rate for generator, sgd for discriminator') artifact = wandb.Artifact('code', type='code') artifact.add_file('train.py') artifact.add_file('losses.py') @@ -46,6 +46,7 @@ class Trainer(object): self.global_batch_size = t_params['batch_size'] self.n_total_image = t_params['n_total_image'] self.enable_unsupervised_losses = t_params['enable_unsupervised_losses'] + self.disable_l1_loss = t_params['disable_l1_loss'] self.max_steps = int(np.ceil(self.n_total_image / self.global_batch_size)) self.n_samples = min(t_params['batch_size'], t_params['n_samples']) self.train_res = t_params['train_res'] @@ -53,6 +54,7 @@ class Trainer(object): self.save_step = 5000 self.image_summary_step = 500 self.unsupervised_coeff = 1 + self.gt_loss_coeff = 1 self.reached_max_steps = False self.log_template = '{:s}, {:s}, {:s}'.format( 'step {}: elapsed: {:.2f}s, d_loss: {:.3f}, g_loss: {:.3f}', @@ -86,6 +88,7 @@ class Trainer(object): beta_1=self.d_opt['beta1'], beta_2=self.d_opt['beta2'], epsilon=self.d_opt['epsilon']) + # self.d_optimizer = tf.keras.optimizers.SGD() self.g_optimizer = tf.keras.optimizers.Adam(self.g_opt['learning_rate'], beta_1=self.g_opt['beta1'], beta_2=self.g_opt['beta2'], @@ -199,6 +202,7 @@ class Trainer(object): def g_train_step_extra(self, dist_inputs): real_images, conditions, mesh_static, mesh_body_posed, mesh_mask = dist_inputs[0] + uv_mask = conditions[4] with tf.GradientTape() as g_tape: # compute losses @@ -207,23 +211,24 @@ class Trainer(object): mesh_pred = uv2mesh(fake_images) e_loss = edge_loss(mesh_pred, mesh_static, mesh_mask) c_loss = collision_penalty(mesh_pred, mesh_body_posed, mesh_mask) - m_loss = 0.0 # uv_mask_loss(fake_images, uv_mask) + l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(real_images - fake_images)), axis=1) * uv_mask) # scale loss g_loss = tf.reduce_sum(g_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 - m_loss = tf.reduce_sum(m_loss) * self.global_batch_scaler + l1_loss = tf.reduce_sum(l1_loss) * self.global_batch_scaler - g_loss = g_loss + self.unsupervised_coeff * (e_loss + c_loss + m_loss) + g_loss = g_loss + self.unsupervised_coeff * (e_loss + c_loss) + self.gt_loss_coeff * l1_loss g_gradients = g_tape.gradient(g_loss, self.generator.trainable_variables) self.g_optimizer.apply_gradients(zip(g_gradients, self.generator.trainable_variables)) - return g_loss, e_loss, c_loss, m_loss + return g_loss, e_loss, c_loss, l1_loss def g_train_step_reg_extra(self, dist_inputs): # real_images, conditions, mesh_static, mesh_body_posed, mesh_mask, uv_mask = dist_inputs[0] real_images, conditions, mesh_static, mesh_body_posed, mesh_mask = dist_inputs[0] + uv_mask = conditions[4] with tf.GradientTape() as g_tape: # compute losses @@ -236,21 +241,21 @@ class Trainer(object): e_loss = edge_loss(mesh_pred, mesh_static, mesh_mask) c_loss = collision_penalty(mesh_pred, mesh_body_posed, mesh_mask) - m_loss = 0.0 # uv_mask_loss(fake_images, uv_mask) + l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(real_images - fake_images)), axis=1) * uv_mask) # scale loss pl_penalty = tf.reduce_sum(pl_penalty) * self.global_batch_scaler g_gan_loss = tf.reduce_sum(g_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 - m_loss = tf.reduce_sum(m_loss) * self.global_batch_scaler + l1_loss = tf.reduce_sum(l1_loss) * self.global_batch_scaler # combine - g_loss = g_gan_loss + pl_penalty + self.unsupervised_coeff * (e_loss + c_loss + m_loss) + g_loss = g_loss + self.unsupervised_coeff * (e_loss + c_loss) + self.gt_loss_coeff * l1_loss g_gradients = g_tape.gradient(g_loss, self.generator.trainable_variables) self.g_optimizer.apply_gradients(zip(g_gradients, self.generator.trainable_variables)) - return g_loss, g_gan_loss, pl_penalty, e_loss, c_loss, m_loss + return g_loss, g_gan_loss, pl_penalty, e_loss, c_loss, l1_loss def train(self, dist_datasets, strategy): @tf.function @@ -287,8 +292,8 @@ class Trainer(object): mean_g_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) - mean_m_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[3], axis=None) - return mean_g_loss, mean_e_loss, mean_c_loss, mean_m_loss + mean_l1_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[3], axis=None) + return mean_g_loss, mean_e_loss, mean_c_loss, mean_l1_loss @tf.function def dist_g_train_step_reg_extra(inputs): @@ -298,8 +303,8 @@ class Trainer(object): mean_pl_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) - mean_m_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[5], axis=None) - return mean_g_loss, mean_g_gan_loss, mean_pl_penalty, mean_e_loss, mean_c_loss, mean_m_loss + mean_l1_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[5], axis=None) + return mean_g_loss, mean_g_gan_loss, mean_pl_penalty, mean_e_loss, mean_c_loss, mean_l1_loss @tf.function def dist_gen_samples(dist_inputs): @@ -326,7 +331,7 @@ class Trainer(object): metric_pl_penalty = tf.keras.metrics.Mean('pl_penalty', dtype=tf.float32) metric_e_loss = tf.keras.metrics.Mean('edge_loss', dtype=tf.float32) metric_c_loss = tf.keras.metrics.Mean('collision_loss', dtype=tf.float32) - metric_m_loss = tf.keras.metrics.Mean('uv_mask_loss', dtype=tf.float32) + metric_l1_loss = tf.keras.metrics.Mean('l1_loss', dtype=tf.float32) # start training zero = tf.constant(0.0, dtype=tf.float32) @@ -346,20 +351,20 @@ class Trainer(object): # g train step if (step + 1) % self.g_opt['reg_interval'] == 0: if step > self.enable_unsupervised_losses: - g_loss, g_gan_loss, pl_penalty, e_loss, c_loss, m_loss = dist_g_train_step_reg_extra((batch, )) + g_loss, g_gan_loss, pl_penalty, e_loss, c_loss, l1_loss = dist_g_train_step_reg_extra((batch, )) else: g_loss, g_gan_loss, pl_penalty = dist_g_train_step_reg((batch, )) e_loss = 0 c_loss = 0 - m_loss = 0 + l1_loss = 0 else: if step > self.enable_unsupervised_losses: - g_loss, e_loss, c_loss, m_loss = dist_g_train_step_extra((batch,)) + g_loss, e_loss, c_loss, l1_loss = dist_g_train_step_extra((batch,)) else: g_loss = dist_g_train_step((batch,)) e_loss = 0 c_loss = 0 - m_loss = 0 + l1_loss = 0 g_gan_loss = g_loss pl_penalty = zero @@ -367,6 +372,11 @@ class Trainer(object): # if step > self.enable_unsupervised_losses: # self.unsupervised_coeff = 1 + + if step < self.disable_l1_loss: + self.gt_loss_coeff = 1 - 1/step + else: + self.gt_loss_coeff = 0 # update g_clone self.g_clone.set_as_moving_average_of(self.generator) @@ -380,7 +390,7 @@ class Trainer(object): metric_pl_penalty(pl_penalty) metric_e_loss(e_loss) metric_c_loss(c_loss) - metric_m_loss(m_loss) + metric_l1_loss(l1_loss) # get current step step += 1 @@ -396,7 +406,8 @@ class Trainer(object): tf.summary.scalar('edge_loss', metric_e_loss.result(), step=step) tf.summary.scalar('collision_loss', metric_c_loss.result(), step=step) tf.summary.scalar('unsupervised_coeff', self.unsupervised_coeff, step=step) - tf.summary.scalar('uv_mask_loss', metric_m_loss.result(), step=step) + tf.summary.scalar('gt_loss_coeff', self.gt_loss_coeff, step=step) + tf.summary.scalar('l1_loss', metric_l1_loss.result(), step=step) # save every self.save_step if step % self.save_step == 0: @@ -417,8 +428,9 @@ class Trainer(object): tf.summary.image('images', summary_image, step=step) template = uv2mesh(batch[1][2]) + mesh_gt = uv2mesh(batch[0]) # gt = uv2mesh(batch[0]) - plot_meshes(mesh_writer, garment=mesh, template=template, body=batch[3], mask=batch[4].numpy(), step=step) + plot_meshes(mesh_writer, garment=mesh, mesh_gt=mesh_gt, template=template, body=batch[3], mask=batch[4].numpy(), step=step) # plot_meshes(mesh_writer, garment=gt, template=template, body=batch[3], mask=batch[4].numpy(), step=step) # train_summary_writer.flush() @@ -472,11 +484,14 @@ def filter_resolutions_featuremaps(resolutions, featuremaps, res): return filtered_resolutions, filtered_featuremaps -def plot_meshes(writer, garment, template, body, mask, step, nr_offset=0): +def plot_meshes(writer, garment, mesh_gt, template, body, mask, step, nr_offset=0): + body = body.numpy() + template = template.numpy() + mesh_gt = mesh_gt.numpy() 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() + template = [template[i, mask.astype(bool)[i,:], :] for i in range(len(garment))] + mesh_gt = [mesh_gt[i, mask.astype(bool)[i,:], :] for i in range(len(garment))] collision = [np.concatenate((body[i], garment[i]), axis=0) for i in range(len(garment))] collision_faces = [np.concatenate((F, faces[i] + 14475), axis=0) for i in range(len(garment))] @@ -488,7 +503,10 @@ def plot_meshes(writer, garment, template, body, mask, step, nr_offset=0): 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[i][np.newaxis,:,:], global_step=step, colors=collision_color[i][np.newaxis,:,:]) - + + for i in range(len(garment)): + plotly_fig = plot_3d(mesh_gt[i], garment[i], faces[i], collision[i], collision_faces[i], collision_color[i]) + wandb.log({"garment_" + str(i): plotly_fig}) def main(): # global program arguments parser @@ -515,7 +533,7 @@ def main(): # network params resolutions = [4, 8, 16, 32, 64, 128, 256, 512, 1024] - g_featuremaps = [256, 256, 256, 256, 256, 128, 64, 32, 16] + g_featuremaps = [128, 128, 128, 128, 128, 64, 32, 16, 8] d_featuremaps = [512, 512, 512, 512, 512, 256, 128, 64, 32] # featuremaps = [512, 512, 512, 512, 512, 256, 128, 64, 32] # featuremaps = [64, 64, 64, 64, 64, 32, 16, 8, 4] @@ -559,13 +577,14 @@ def main(): 'use_custom_cuda': args['use_custom_cuda'], 'model_base_dir': args['model_base_dir'], 'enable_unsupervised_losses': 0, + 'disable_l1_loss': 10000, # network params 'g_params': g_params, 'd_params': d_params, # training params - 'g_opt': {'learning_rate': 0.002, 'beta1': 0.0, 'beta2': 0.99, 'epsilon': 1e-08, 'reg_interval': 8}, + 'g_opt': {'learning_rate': 0.0002, 'beta1': 0.0, 'beta2': 0.99, 'epsilon': 1e-08, 'reg_interval': 8}, 'd_opt': {'learning_rate': 0.002, 'beta1': 0.0, 'beta2': 0.99, 'epsilon': 1e-08, 'reg_interval': 16}, 'batch_size': global_batch_size, 'n_total_image': 25000000, @@ -573,7 +592,7 @@ def main(): 'train_res': args['train_res'], } - trainer = Trainer(training_parameters, name=f'stylegan2-test-16-{args["train_res"]}x{args["train_res"]}') + trainer = Trainer(training_parameters, name=f'stylegan2-test-19-{args["train_res"]}x{args["train_res"]}') trainer.train(dist_dataset, strategy) return