test-20

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):
@@ -86,3 +87,52 @@ def maskFacesBatch(V, F, mask):
         Vs.append(v)
         Fs.append(f)
     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

requierements.txt

@@ -5,3 +5,4 @@ wandb
 OpenEXR
 tensorflow-graphics
 scipy
+plotly

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)
+            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_m_loss
+            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)
+            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_m_loss
+            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
@@ -368,6 +373,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))
+    template = [template[i, mask.astype(bool)[i,:], :] for i in range(len(garment))]
-    body = body.numpy()
+    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))]
@@ -489,6 +504,9 @@ def plot_meshes(writer, garment, template, body, mask, step, nr_offset=0):
         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