updated pipeline, added uv_mask_loss

custom_utils/custom_utils.py

@@ -21,7 +21,7 @@ class UV2Mesh():
         super(UV2Mesh, self).__init__()
         self.res = res
         uv_map = np.load('custom_utils/stretched_uv_map_3.npy')
-        self.uv_map_discreet = np.round(uv_map * self.res).astype(np.int32)
+        self.uv_map_discreet = np.floor(uv_map * self.res).astype(np.int32)
         m2uv = np.load('custom_utils/m2uv.npy', allow_pickle=True).item()
         m2uv = m2uv_array(m2uv)
         self.uv_pixels = self.uv_map_discreet[m2uv]
@@ -35,6 +35,7 @@ class UV2Mesh():
         # self.z2 = tf.concat((tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [self.batch_size, 1, 1]), tf.ones((self.batch_size, 14475, 1), dtype=tf.int32) * 2), axis=-1)
 
     # input = channels_first
+    @tf.function
     def uv_to_mesh(self, uv):
         uv = tf.transpose(uv, [0,2,3,1])
         batch_size = uv.shape[0]

dataset.py

@@ -23,15 +23,16 @@ class DatasetSimple(tf.data.Dataset):
             uv_static = np.load(os.path.join(folder, 'uv_static', id +  '.npz')); uv_static = uv_static['uv_static']; uv_static = np.transpose(uv_static, [2, 0, 1])
             mesh_static = np.load(os.path.join(folder, 'mesh_static', id +  '.npz')); mesh_static = mesh_static['mesh_static']
             mesh_body_posed = np.load(os.path.join(folder, 'mesh_body_posed', id + '_' + frame_nr + '.npy'))
-            mask = np.load(os.path.join(folder, 'mask', id + '.npy')); mask = mask.astype(np.float32)
+            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)
 
             # uv_posed, (pose, shape, uv_static), mesh_static, mesh_body_posed, mask
-            yield uv_posed, (pose_shape[:72], pose_shape[72:], uv_static), mesh_static, mesh_body_posed, mask
+            yield uv_posed, (pose_shape[:72], pose_shape[72:], uv_static), mesh_static, mesh_body_posed, mesh_mask, uv_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 = (72), dtype = tf.float32), tf.TensorSpec(shape = (11), dtype = tf.float32), tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32) ), tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32), tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32), tf.TensorSpec(shape = (14475,), dtype = tf.float32) ),
+            output_signature = (tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32), (tf.TensorSpec(shape = (72), dtype = tf.float32), tf.TensorSpec(shape = (11), dtype = tf.float32), tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32) ), tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32), tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32), tf.TensorSpec(shape = (14475,), dtype = tf.float32), tf.TensorSpec(shape = (res * 2, res,), dtype = tf.float32) ),
             args=(folder, )
         )
 

inference_custom.py

@@ -9,6 +9,7 @@ from dataset import DatasetSimple
 from load_models import load_generator
 from tf_utils import allow_memory_growth
 
+os.environ["CUDA_VISIBLE_DEVICES"] = "3"
 
 def m2uv_array(m2uv):
     result = np.zeros((14475, 2), dtype=np.int32)
@@ -39,7 +40,7 @@ def uv_to_mesh(uv):
     return mesh
 
 
-def inference(ckpt_dir, use_custom_cuda, res, folder, data_folder, samples=100):
+def inference(ckpt_dir, use_custom_cuda, res, folder, data_folder, samples=50):
     # create generator
     resolutions = [4, 8, 16, 32, 64, 128, 256, 512, 1024]
     featuremaps = [512, 512, 512, 512, 512, 256, 128, 64, 32]
@@ -53,7 +54,9 @@ def inference(ckpt_dir, use_custom_cuda, res, folder, data_folder, samples=100):
         '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()
+    dataset = DatasetSimple(data_folder).cache().batch(1).repeat()
+
+    os.makedirs(os.path.join(folder), exist_ok=True)
 
     i = 0
     for batch in dataset:
@@ -80,9 +83,9 @@ def main():
     checkpoints = [
         {
             'res': 512,
-            'ckpt_dir': 'models/stylegan2-test-5-512x512',
+            'ckpt_dir': 'models/stylegan2-test-8-512x512',
             'use_custom_cuda': True,
-            'folder': 'output',
+            'folder': 'output/test-8-196000',
             'data_folder': '/home/workspace/data/dataset/stylegan_new/train'
         },
     ]

losses.py

@@ -187,6 +187,17 @@ def collision_penalty(va, vb, mask, eps=1e-3):
 
     return tf.reduce_sum(interpenetration**3) / batch_size
 
+whole_garment_uv_mask = np.load('custom_utils/mask_256x512.npy')
+
+def uv_mask_loss(pred, uv_mask):
+    inv_mask = (1 - uv_mask)  * whole_garment_uv_mask
+    pred = tf.reduce_sum(pred, axis=1)
+    invalid = pred * inv_mask
+    invalid = invalid > 0.01
+    invalid_count = tf.math.count_nonzero(invalid, dtype=tf.float32)
+    total_count = tf.math.count_nonzero(inv_mask, dtype=tf.float32)
+    return invalid_count / total_count
+
 
 if __name__ == '__main__':
     V = tf.ones((4, 14475, 3))

train.py

@@ -12,14 +12,17 @@ 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 import DatasetSimple
-from losses import d_logistic, d_logistic_r1_reg, g_logistic_non_saturating, g_logistic_ns_pathreg, edge_loss
+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
 
 
-os.environ["CUDA_VISIBLE_DEVICES"] = "1"
+os.environ["CUDA_VISIBLE_DEVICES"] = "2"
-wandb.init(config=tf.compat.v1.flags.FLAGS, sync_tensorboard=True, project='stylegan', name='test_conditioned_2')
+run = wandb.init(config=tf.compat.v1.flags.FLAGS, sync_tensorboard=True, project='stylegan', name='test_conditioned_7(test-10)')
-
+artifact = wandb.Artifact('code', type='code')
+artifact.add_file('train.py')
+artifact.add_file('losses.py')
+run.log_artifact(artifact)
 
 def initiate_models(g_params, d_params, use_custom_cuda):
     discriminator = load_discriminator(d_params, ckpt_dir=None, custom_cuda=use_custom_cuda)
@@ -41,18 +44,19 @@ class Trainer(object):
         self.model_base_dir = t_params['model_base_dir']
         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.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']
         self.print_step = 10
-        self.save_step = 1000
+        self.save_step = 5000
         self.image_summary_step = 500
+        self.unsupervised_coeff = 1
         self.reached_max_steps = False
         self.log_template = '{:s}, {:s}, {:s}'.format(
             'step {}: elapsed: {:.2f}s, d_loss: {:.3f}, g_loss: {:.3f}',
             'd_gan_loss: {:.3f}, g_gan_loss: {:.3f}',
-            'r1_penalty: {:.3f}, pl_penalty: {:.3f}',
+            'r1_penalty: {:.3f}, pl_penalty: {:.3f}')
-            'e_loss: {:.3f}, c_loss: {:.3f}')
 
         # copy network params
         self.g_params = t_params['g_params']
@@ -94,7 +98,7 @@ class Trainer(object):
                                         generator=self.generator,
                                         g_clone=self.g_clone,
                                         pl_mean=self.pl_mean)
-        self.manager = tf.train.CheckpointManager(self.ckpt, self.ckpt_dir, max_to_keep=5)
+        self.manager = tf.train.CheckpointManager(self.ckpt, self.ckpt_dir, max_to_keep=20)
 
         # try to restore
         self.ckpt.restore(self.manager.latest_checkpoint)
@@ -120,30 +124,24 @@ class Trainer(object):
         return params_copy
 
     def d_train_step(self, dist_inputs):
-        real_images, conditions, _, _, _ = dist_inputs[0]
+        real_images, conditions, _, _, _, _ = dist_inputs[0]
 
         with tf.GradientTape() as d_tape:
             # compute losses
             d_loss = d_logistic(real_images, self.generator, self.discriminator, self.g_params['z_dim'], conditions=conditions)
 
-            # mesh_pred = uv2mesh(fake_images)
-
-            # 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
 
-            d_loss = d_loss# + e_loss + c_loss
+            d_loss = d_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
 
     def d_train_step_reg(self, dist_inputs):
-        real_images, conditions, _, _, _ = dist_inputs[0]
+        real_images, conditions, _, _, _, _ = dist_inputs[0]
 
         with tf.GradientTape() as d_tape:
             # compute losses
@@ -151,49 +149,80 @@ class Trainer(object):
                                                        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)
-
-            # 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
 
             # combine
-            d_loss = d_gan_loss + r1_penalty# + e_loss + c_loss
+            d_loss = d_gan_loss + r1_penalty
 
         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
+
 
     def g_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 g_tape:
+            # compute losses
+            _, g_loss = g_logistic_non_saturating(real_images, self.generator, self.discriminator, self.g_params['z_dim'], conditions=conditions)
+
+            # scale loss
+            g_loss = tf.reduce_sum(g_loss) * self.global_batch_scaler
+            
+        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
+
+    def g_train_step_reg(self, dist_inputs):
+        real_images, conditions, _, _, _, _ = dist_inputs[0]
+
+        with tf.GradientTape() as g_tape:
+            # compute losses
+            _, g_gan_loss, pl_penalty = g_logistic_ns_pathreg(real_images, self.generator, self.discriminator,
+                                                           self.g_params['z_dim'], self.pl_mean,
+                                                           self.pl_minibatch_shrink, self.pl_denorm, self.pl_decay, conditions=conditions)
+            pl_penalty = pl_penalty * self.pl_weight * self.g_opt['reg_interval']
+
+            # 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
+            
+            # combine
+            g_loss = g_gan_loss + pl_penalty
+
+        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
+
+
+    def g_train_step_extra(self, dist_inputs):
+        real_images, conditions, mesh_static, mesh_body_posed, mesh_mask, uv_mask = dist_inputs[0]
 
         with tf.GradientTape() as g_tape:
             # compute losses
             fake_images, g_loss = g_logistic_non_saturating(real_images, self.generator, self.discriminator, self.g_params['z_dim'], conditions=conditions)
 
             mesh_pred = uv2mesh(fake_images)
-
+            e_loss = 0.1 * edge_loss(mesh_pred, mesh_static, mesh_mask)
-            e_loss = edge_loss(mesh_pred, mesh_static, mask)
+            c_loss = 0.0 # collision_penalty(mesh_pred, mesh_body_posed, mesh_mask)
-            c_loss = collision_penalty(mesh_pred, mesh_body_posed, mask)
+            m_loss = uv_mask_loss(fake_images, 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
 
-            g_loss = g_loss + e_loss + c_loss
+            g_loss = g_loss + self.unsupervised_coeff * (e_loss + c_loss + m_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
+        return g_loss, e_loss, c_loss, m_loss
 
-    def g_train_step_reg(self, dist_inputs):
+    def g_train_step_reg_extra(self, dist_inputs):
-        real_images, conditions, mesh_static, mesh_body_posed, mask = dist_inputs[0]
+        real_images, conditions, mesh_static, mesh_body_posed, mesh_mask, uv_mask = dist_inputs[0]
 
         with tf.GradientTape() as g_tape:
             # compute losses
@@ -204,30 +233,30 @@ class Trainer(object):
 
             mesh_pred = uv2mesh(fake_images)
 
-            e_loss = edge_loss(mesh_pred, mesh_static, mask)
+            e_loss = 0.1 * edge_loss(mesh_pred, mesh_static, mesh_mask)
-            c_loss = collision_penalty(mesh_pred, mesh_body_posed, mask)
+            c_loss = 0.0 # collision_penalty(mesh_pred, mesh_body_posed, mesh_mask)
+            m_loss = uv_mask_loss(fake_images, 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
             
             # combine
-            g_loss = g_gan_loss + pl_penalty + e_loss + c_loss
+            g_loss = g_gan_loss + pl_penalty + self.unsupervised_coeff * (e_loss + c_loss + m_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
+        return g_loss, g_gan_loss, pl_penalty, e_loss, c_loss, m_loss
 
     def train(self, dist_datasets, strategy):
         @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, axis=None)
-            # mean_e_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[1], axis=None)
+            return mean_d_loss
-            # 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,17 +264,13 @@ 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)
+            return mean_d_loss, mean_d_gan_loss, mean_r1_penalty
-            # 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):
             per_replica_losses = strategy.run(fn=self.g_train_step, args=(inputs,))
-            mean_g_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[0], axis=None)
+            mean_g_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)
+            return mean_g_loss
-            mean_c_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[2], axis=None)
-            return mean_g_loss, mean_e_loss, mean_c_loss
 
         @tf.function
         def dist_g_train_step_reg(inputs):
@@ -253,9 +278,27 @@ class Trainer(object):
             mean_g_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[0], axis=None)
             mean_g_gan_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[1], axis=None)
             mean_pl_penalty = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[2], axis=None)
+            return mean_g_loss, mean_g_gan_loss, mean_pl_penalty
+
+        @tf.function
+        def dist_g_train_step_extra(inputs):
+            per_replica_losses = strategy.run(fn=self.g_train_step_extra, args=(inputs,))
+            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
+
+        @tf.function
+        def dist_g_train_step_reg_extra(inputs):
+            per_replica_losses = strategy.run(fn=self.g_train_step_reg_extra, args=(inputs,))
+            mean_g_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[0], axis=None)
+            mean_g_gan_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[1], axis=None)
+            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)
-            return mean_g_loss, mean_g_gan_loss, mean_pl_penalty, mean_e_loss, mean_c_loss
+            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
 
         @tf.function
         def dist_gen_samples(dist_inputs):
@@ -282,6 +325,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)
 
         # start training
         zero = tf.constant(0.0, dtype=tf.float32)
@@ -300,12 +344,29 @@ class Trainer(object):
 
             # g train step
             if (step + 1) % self.g_opt['reg_interval'] == 0:
-                g_loss, g_gan_loss, pl_penalty, e_loss, c_loss = dist_g_train_step_reg((batch, ))
+                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, ))
                 else:
-                g_loss, e_loss, c_loss = dist_g_train_step((batch,))
+                    g_loss, g_gan_loss, pl_penalty = dist_g_train_step_reg((batch, ))
+                    e_loss = 0
+                    c_loss = 0
+                    m_loss = 0
+            else:
+                if step > self.enable_unsupervised_losses:
+                    g_loss, e_loss, c_loss, m_loss = dist_g_train_step_extra((batch,))
+                else:
+                    g_loss = dist_g_train_step((batch,))
+                    e_loss = 0
+                    c_loss = 0
+                    m_loss = 0
+
                 g_gan_loss = g_loss
                 pl_penalty = zero
                 
+
+            # if step > self.enable_unsupervised_losses:
+            #     self.unsupervised_coeff = (step / 40000)
+            
             # update g_clone
             self.g_clone.set_as_moving_average_of(self.generator)
 
@@ -318,6 +379,7 @@ class Trainer(object):
             metric_pl_penalty(pl_penalty)
             metric_e_loss(e_loss)
             metric_c_loss(c_loss)
+            metric_m_loss(m_loss)
 
             # get current step
             step += 1
@@ -332,6 +394,8 @@ class Trainer(object):
                 tf.summary.scalar('pl_penalty', metric_pl_penalty.result(), step=step)
                 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)
 
             # save every self.save_step
             if step % self.save_step == 0:
@@ -344,14 +408,17 @@ class Trainer(object):
                 test_conditions = batch[1] #tf.ones((self.n_samples, self.g_params['labels_dim']), dtype=tf.dtypes.float32)
                 summary_image = dist_gen_samples((test_z, test_conditions))
 
+                mesh = uv2mesh(summary_image[:, :, :self.train_res, int(self.train_res/2):])
                 # convert to tensor image
                 summary_image = self.convert_per_replica_image(summary_image, strategy)
 
                 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]))
+                template = uv2mesh(batch[1][2])
-                # plot_meshes(mesh_writer, garment=mesh, template=batch[1][2], body=batch[3], mask=batch[4], step=step)
+                # 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=gt, template=template, body=batch[3], mask=batch[4].numpy(), step=step)
 
                 # train_summary_writer.flush()
 
@@ -360,8 +427,7 @@ class Trainer(object):
                 elapsed = time.time() - t_start
                 print(self.log_template.format(step, elapsed, d_loss.numpy(), g_loss.numpy(),
                                                d_gan_loss.numpy(), g_gan_loss.numpy(),
-                                               r1_penalty.numpy(), pl_penalty.numpy(),
+                                               r1_penalty.numpy(), pl_penalty.numpy()))
-                                               e_loss.numpy(), c_loss.numpy()))
 
                 # reset timer
                 t_start = time.time()
@@ -376,14 +442,17 @@ class Trainer(object):
         return
 
     def gen_samples(self, inputs):
-        test_z, test_labels = inputs
+        test_z, test_conditions = inputs
 
         # run networks
-        fake_images_05 = self.g_clone([test_z, test_labels], truncation_psi=0.5, training=False)
+        fake_images_05 = self.g_clone([test_z, test_conditions], truncation_psi=0.5, training=False)
-        fake_images_07 = self.g_clone([test_z, test_labels], truncation_psi=0.7, training=False)
+        fake_images_07 = self.g_clone([test_z, test_conditions], truncation_psi=0.7, training=False)
 
         # merge on batch dimension: [n_samples, 3, out_res, 2 * out_res]
-        final_image = tf.concat([fake_images_05, fake_images_07], axis=2)
+        template = tf.concat([test_conditions[2], test_conditions[2]], axis=2)
+        fake_image = tf.concat([fake_images_05, fake_images_07], axis=2)
+        final_image = tf.concat([template, fake_image], axis=3)
+
         return final_image
 
     @staticmethod
@@ -409,14 +478,15 @@ def plot_meshes(writer, garment, template, body, mask, step, nr_offset=0):
     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))]
+    collision_faces = [np.concatenate((F, faces[i] + 14475), axis=0) for i in range(len(garment))]
+    collision_color = [np.concatenate((np.ones_like(body[i], dtype=np.int32)*100, np.ones_like(garment[i], dtype=np.int32)*230), 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)
+        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,:,:])
         
 
 def main():
@@ -444,7 +514,8 @@ def main():
 
     # network params
     resolutions = [4, 8, 16, 32, 64, 128, 256, 512, 1024]
-    featuremaps = [512, 512, 512, 512, 512, 256, 128, 64, 32]
+    featuremaps = [256, 256, 256, 256, 256, 128, 64, 32, 16]
+    # featuremaps = [512, 512, 512, 512, 512, 256, 128, 64, 32]
     # featuremaps = [64, 64, 64, 64, 64, 32, 16, 8, 4]
     train_resolutions, train_featuremaps = filter_resolutions_featuremaps(resolutions, featuremaps, args['train_res'])
     g_params = {
@@ -468,12 +539,6 @@ def main():
     # prepare dataset
     # dataset = get_ffhq_dataset(args['tfrecord_dir'], args['train_res'], batch_size=global_batch_size, epochs=None)
     dataset = DatasetSimple('/home/workspace/data/dataset/stylegan_new/train').cache().batch(global_batch_size).repeat()
-    # limit = int(25594/global_batch_size)
-    # count = 0
-    # for d in dataset:
-    #     count += global_batch_size
-    #     if count > limit:
-    #         break
 
     options = tf.data.Options()
     options.experimental_distribute.auto_shard_policy = tf.data.experimental.AutoShardPolicy.OFF
@@ -490,6 +555,7 @@ def main():
             'use_tf_function': args['use_tf_function'],
             'use_custom_cuda': args['use_custom_cuda'],
             'model_base_dir': args['model_base_dir'],
+            'enable_unsupervised_losses': 0,
 
             # network params
             'g_params': g_params,
@@ -504,7 +570,7 @@ def main():
             'train_res': args['train_res'],
         }
 
-        trainer = Trainer(training_parameters, name=f'stylegan2-test-5-{args["train_res"]}x{args["train_res"]}')
+        trainer = Trainer(training_parameters, name=f'stylegan2-test-10-{args["train_res"]}x{args["train_res"]}')
         trainer.train(dist_dataset, strategy)
     return