test-16

dataset.py

@@ -19,6 +19,7 @@ class DatasetSimple(tf.data.Dataset):
         for frame in frames:
             id, frame_nr = frame.split('.')[0].split('_')
             uv_posed = np.load(os.path.join(folder, 'uv_posed', id + '_' + frame_nr + '.npz')); uv_posed = uv_posed['uv_gt']; uv_posed = np.transpose(uv_posed, [2, 0, 1])
+            uv_body_posed = np.load(os.path.join(folder, 'uv_body_posed', id + '_' + frame_nr + '.npz')); uv_body_posed = uv_body_posed['uv_body']; uv_body_posed = np.transpose(uv_body_posed, [2, 0, 1])
             pose_shape = np.load(os.path.join(folder, 'pose_shape',  id + '_' + frame_nr + '.npy'))
             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']
@@ -27,12 +28,12 @@ class DatasetSimple(tf.data.Dataset):
             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, uv_mask), mesh_static, mesh_body_posed, mesh_mask
+            yield uv_posed, (pose_shape[:72], pose_shape[72:], uv_static, uv_mask, uv_body_posed), mesh_static, mesh_body_posed, mesh_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 = (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 = (res * 2, res,), 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) ),
             args=(folder, )
         )
 

inference_custom.py

@@ -9,7 +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"
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
 
 def m2uv_array(m2uv):
     result = np.zeros((14475, 2), dtype=np.int32)
@@ -83,9 +83,9 @@ def main():
     checkpoints = [
         {
             'res': 512,
-            'ckpt_dir': 'models/stylegan2-test-12-512x512',
+            'ckpt_dir': 'models/stylegan2-test-15-512x512',
             'use_custom_cuda': True,
-            'folder': 'output/test-12-140000',
+            'folder': 'output/test-15-115000',
             'data_folder': '/home/workspace/data/dataset/stylegan_new/train'
         },
     ]

load_models.py

@@ -18,7 +18,7 @@ def load_generator(g_params=None, is_g_clone=False, ckpt_dir=None, custom_cuda=T
         }
 
     test_latent = tf.ones((1, g_params['z_dim']), dtype=tf.float32)
-    test_conditions = (tf.ones((1, 72), dtype=tf.float32), tf.ones((1, 11), dtype=tf.float32), tf.ones((1, 3, 512, 256), dtype=tf.float32), tf.ones((1, 512, 256), dtype=tf.float32))
+    test_conditions = (tf.ones((1, 72), dtype=tf.float32), tf.ones((1, 11), dtype=tf.float32), tf.ones((1, 3, 512, 256), dtype=tf.float32), tf.ones((1, 512, 256), dtype=tf.float32), tf.ones((1, 3, 512, 256), dtype=tf.float32))
     # test_conditions = (tf.ones((1, 72), dtype=tf.float32), tf.ones((1, 11), dtype=tf.float32), tf.ones((1, 3, 512, 256), dtype=tf.float32))
 
     # build generator model

losses.py

@@ -70,7 +70,7 @@ def g_logistic_ns_pathreg(real_images, generator, discriminator, z_dim,
 
     pl_minibatch = tf.maximum(1, tf.math.floordiv(batch_size, pl_minibatch_shrink))
     pl_z = tf.random.normal(shape=[pl_minibatch, z_dim], dtype=tf.float32)
-    pl_conditions = (conditions[0][:pl_minibatch], conditions[1][:pl_minibatch], conditions[2][:pl_minibatch], conditions[3][:pl_minibatch])
+    pl_conditions = (conditions[0][:pl_minibatch], conditions[1][:pl_minibatch], conditions[2][:pl_minibatch], conditions[3][:pl_minibatch], conditions[4][:pl_minibatch])
 
     # forward pass
     fake_images = generator([z, conditions], ret_w_broadcasted=False, training=True)

stylegan2/generator.py

@@ -102,6 +102,7 @@ class Generator(tf.keras.Model):
             w_broadcasted = self.truncation_trick(w_broadcasted, truncation_psi, truncation_cutoff)
 
         image_out = self.synthesis(w_broadcasted)
+        image_out = image_out + conditions[4]
         image_out = image_out * conditions[3][:,tf.newaxis,:,:]
 
         if ret_w_broadcasted:

train.py

@@ -17,8 +17,8 @@ from losses import d_logistic, d_logistic_r1_reg, g_logistic_non_saturating, g_l
 from losses import collision_loss as collision_penalty
 
 
-os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+os.environ["CUDA_VISIBLE_DEVICES"] = "3"
-run = wandb.init(config=tf.compat.v1.flags.FLAGS, sync_tensorboard=True, project='stylegan', name='test_conditioned_10(test-13)', notes='same as test-12 but double discriminator feature sizes')
+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')
 artifact = wandb.Artifact('code', type='code')
 artifact.add_file('train.py')
 artifact.add_file('losses.py')
@@ -206,7 +206,7 @@ class Trainer(object):
 
             mesh_pred = uv2mesh(fake_images)
             e_loss = edge_loss(mesh_pred, mesh_static, mesh_mask)
-            c_loss = 0.0 #collision_penalty(mesh_pred, mesh_body_posed, mesh_mask)
+            c_loss = collision_penalty(mesh_pred, mesh_body_posed, mesh_mask)
             m_loss = 0.0 # uv_mask_loss(fake_images, uv_mask)
 
             # scale loss
@@ -235,7 +235,7 @@ class Trainer(object):
             mesh_pred = uv2mesh(fake_images)
 
             e_loss = edge_loss(mesh_pred, mesh_static, mesh_mask)
-            c_loss = 0.0 #collision_penalty(mesh_pred, mesh_body_posed, mesh_mask)
+            c_loss = collision_penalty(mesh_pred, mesh_body_posed, mesh_mask)
             m_loss = 0.0 # uv_mask_loss(fake_images, uv_mask)
 
             # scale loss
@@ -366,7 +366,7 @@ class Trainer(object):
                 
 
             # if step > self.enable_unsupervised_losses:
-            #     self.unsupervised_coeff = (step / 40000)
+            #     self.unsupervised_coeff = 1
             
             # update g_clone
             self.g_clone.set_as_moving_average_of(self.generator)
@@ -573,7 +573,7 @@ def main():
             'train_res': args['train_res'],
         }
 
-        trainer = Trainer(training_parameters, name=f'stylegan2-test-13-{args["train_res"]}x{args["train_res"]}')
+        trainer = Trainer(training_parameters, name=f'stylegan2-test-16-{args["train_res"]}x{args["train_res"]}')
         trainer.train(dist_dataset, strategy)
     return