plotting not working, losses fixed?

.gitignore

@@ -135,3 +135,5 @@ dmypy.json
 .vscode/
 
 wandb/
+output/
+output.zip

custom_utils/custom_utils.py

@@ -63,3 +63,25 @@ class UV2Mesh():
     def __call__(self, uv):
         return self.uv_to_mesh(uv)
         
+
+
+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

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',
+            'ckpt_dir': 'models/stylegan2-test-5-512x512',
-            'use_custom_cuda': False,
+            'use_custom_cuda': True,
-            'folder': '/workspace/Checkpoints/checkpoint_52000/output',
+            '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
 
 

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]))
+    m = tf.ones((4, 14475))
-    loss_1 = collision_loss(V, B, thr=.004, stop_gradient=False)
+    loss_1 = collision_loss(V, B, m, thr=.004)
-    loss_2 = collision_penalty(V, B, B_vn)
     print()

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"
+os.environ["CUDA_VISIBLE_DEVICES"] = "1"
-wandb.init(config=tf.compat.v1.flags.FLAGS, sync_tensorboard=True, project='stylegan', name='test_conditioned_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)
+            # 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)
 
             # scale loss
             d_loss = tf.reduce_sum(d_loss) * self.global_batch_scaler
-            e_loss = tf.reduce_sum(e_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
+            # 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)
+            # 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)
 
             # 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
+            # e_loss = tf.reduce_sum(e_loss) * self.global_batch_scaler
-            c_loss = tf.reduce_sum(c_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_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_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_c_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[2], axis=None)
-            return mean_d_loss, mean_e_loss, mean_c_loss
+            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_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_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
+            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