train on template only with mask applied

.gitignore

@@ -135,5 +135,9 @@ dmypy.json
 .vscode/
 
 wandb/
-output/
-output.zip
+output/*
+*.zip
+*.out
+models/
+models_archive/
+temp.jpg

dataset.py

@@ -27,12 +27,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), mesh_static, mesh_body_posed, mesh_mask, uv_mask
+            yield uv_posed, (pose_shape[:72], pose_shape[72:], uv_static, uv_mask), 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 = (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) ),
+            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) ),
             args=(folder, )
         )
 

inference_custom.py

@@ -43,7 +43,7 @@ def uv_to_mesh(uv):
 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]
+    featuremaps = [256, 256, 256, 256, 256, 128, 64, 32, 16]
     filter_index = resolutions.index(res)
     g_params = {
         'z_dim': 256,
@@ -66,7 +66,7 @@ def inference(ckpt_dir, use_custom_cuda, res, folder, data_folder, samples=50):
                                 batch[1]],
                                 training=False, truncation_psi=0.5)
 
-        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]))
+        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]), gt=tf.squeeze(batch[0]))
         img = np.transpose(np.round((fake_images + 1) * 128)[0], [1,2,0])
         cv2.imwrite(os.path.join(folder, str(i) + '.jpg'), img)
         i += 1
@@ -83,9 +83,9 @@ def main():
     checkpoints = [
         {
             'res': 512,
-            'ckpt_dir': 'models/stylegan2-test-8-512x512',
+            'ckpt_dir': 'models/stylegan2-test-12-512x512',
             'use_custom_cuda': True,
-            'folder': 'output/test-8-196000',
+            'folder': 'output/test-12-140000',
             'data_folder': '/home/workspace/data/dataset/stylegan_new/train'
         },
     ]

load_models.py

@@ -18,7 +18,8 @@ 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))
+    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))
 
     # build generator model
     generator = Generator(g_params)

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])
+    pl_conditions = (conditions[0][:pl_minibatch], conditions[1][:pl_minibatch], conditions[2][:pl_minibatch], conditions[3][:pl_minibatch])
 
     # forward pass
     fake_images = generator([z, conditions], ret_w_broadcasted=False, training=True)

models/.gitignore

@@ -1,5 +0,0 @@
-# ignore all
-*
-
-# except this file
-!.gitignore

reconstruct.py

@@ -22,10 +22,10 @@ def m2uv_array(m2uv):
 def uv_to_mesh(uv):
     uv_res = 256
     # m2uv = np.load('/workspace/Code/featurelines/MeshUtils/m2uv.npy', allow_pickle=True).item()
-    m2uv = np.load(r'D:\Work\Code\featurelines\MeshUtils\m2uv.npy', allow_pickle=True).item()
+    m2uv = np.load('custom_utils\m2uv.npy', allow_pickle=True).item()
     m2uv = m2uv_array(m2uv)
     # uv_map = np.load('/workspace/Code/featurelines/Meshutils/stretched_uv_map_3.npy')
-    uv_map = np.load(r'D:\Work\Code\featurelines\MeshUtils\stretched_uv_map_3.npy')
+    uv_map = np.load('custom_utils\stretched_uv_map_3.npy')
     uv_map_discreet = np.floor(uv_map * uv_res).astype(np.int32)
 
     mesh = np.zeros((14475, 3))
@@ -36,36 +36,84 @@ def uv_to_mesh(uv):
     return mesh
 
 
-def plot(uv, mesh):
-    fig = plt.figure(figsize=(15, 15))
-    ax = fig.add_subplot(111, projection='3d')
-    ax.set_xlim3d(-1, 1)
-    ax.set_ylim3d(-1, 1)
-    ax.set_zlim3d(-1, 1)
-    ax.set_xlabel("X")
-    ax.set_ylabel("Y")
-    ax.set_zlabel("Z")
-    ax.scatter(mesh[:, 0], mesh[:, 1], mesh[:, 2], s=5)
-    # plt.show()
+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
+
+def plot(uv, mesh, faces, body, F, template, uv_gt, mesh_gt):
+    # fig = plt.figure(figsize=(15, 15))
+    # ax = fig.add_subplot(111, projection='3d')
+    # ax.set_xlim3d(-1, 1)
+    # ax.set_ylim3d(-1, 1)
+    # ax.set_zlim3d(-1, 1)
+    # ax.set_xlabel("X")
+    # ax.set_ylabel("Y")
+    # ax.set_zlabel("Z")
+    # ax.scatter(mesh[:, 0], mesh[:, 1], mesh[:, 2], s=5)
+    # # plt.show()
 
     fig = plt.figure(figsize=(15, 15))
-    ax = fig.add_subplot(111)
-    ax.imshow(np.round( (uv - uv.min()) / (uv.max() - uv.min()) * 255).astype(np.int32))
+    ax = fig.add_subplot(131)
+    ax.imshow(np.round((uv + 1) * 128).astype(np.int32))
+    ax = fig.add_subplot(132)
+    ax.imshow(np.round((template + 1) * 128).astype(np.int32))
+    ax = fig.add_subplot(133)
+    ax.imshow(np.round((uv_gt + 1) * 128).astype(np.int32))
     plt.show()
 
-    F = np.load(r'D:\Work\Code\featurelines\FeatureLines\superSMPL\smpl_faces.npy')
-    tm.Trimesh(mesh, F).show()
+    collision = np.concatenate((body, mesh), axis=0)
+    collision_faces = np.concatenate((F, faces + 14475), axis=0)
+    collision_color = np.concatenate((np.ones_like(body, dtype=np.int32)*100, np.ones_like(mesh, dtype=np.int32)*230), axis=0)
+
+
+    tm.Trimesh(mesh, faces).show()
+    tm.Trimesh(mesh_gt, faces).show()
+    # tm.Trimesh(collision, collision_faces, vertex_colors=collision_color).show()
 
 
 
 if __name__ == '__main__':
     file = sys.argv[1]
 
-    uv = np.load(file)
-    uv = np.transpose(uv['uv'][0], [1,2,0])
-    mesh = uv_to_mesh(uv)
+    for i in range(50):
+        file = r'Z:\Work\Huni\Barcelona\Output\stylegan\test-12-15000\\' + str(i) + '.npz'
 
-    plot(uv, mesh)
+        F = np.load('custom_utils\smpl_faces.npy')
+        data = np.load(file)
+        uv = np.transpose(data['uv'], [1,2,0])
+        uv_gt = np.transpose(data['gt'], [1,2,0])
+        mask = data['mask'].astype(bool)
+        template = np.transpose(data['template'], [1,2,0])
+
+
+        mesh = uv_to_mesh(uv)
+        mesh_gt = uv_to_mesh(uv_gt)
+
+        template_mesh = uv_to_mesh(template)
+        filter = np.abs(np.sum(mesh, axis=-1)) < 1e-2
+        overlap = np.sum((((np.sum(template_mesh, axis=-1) * mask) * (np.sum(mesh, axis=-1) * filter * mask))) > 0)
+
+        mesh, faces = _maskFaces(mesh, F, mask)
+        plot(uv, mesh, faces, data['body'], F, template, uv_gt, mesh_gt[mask, :])
+
+        pass
 
     pass
 

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[3][:,tf.newaxis,:,:]
 
         if ret_w_broadcasted:
             return image_out, w_broadcasted

stylegan2/layers/mapping_block.py

@@ -14,8 +14,9 @@ class Mapping(tf.keras.layers.Layer):
         self.gain = 1.0
         self.lrmul = 0.01
 
-        self.pose_shape_embedding = PoseShapeEmbedding()
-        self.garment_embedding = GarmentEmbedding((self.w_dim * 2) - (216 + 11))
+        # self.pose_shape_embedding = PoseShapeEmbedding()
+        # self.garment_embedding = GarmentEmbedding((self.w_dim * 2) - (216 + 11))
+        self.garment_embedding = GarmentEmbedding(self.w_dim)
 
         self.normalize = tf.keras.layers.Lambda(lambda x: x * tf.math.rsqrt(tf.reduce_mean(tf.square(x), axis=1, keepdims=True) + 1e-8))
 
@@ -30,9 +31,10 @@ class Mapping(tf.keras.layers.Layer):
         x = latents
 
         # embed label if any
-        pose_shape_latent = self.pose_shape_embedding([conditions[0], conditions[1]])
+        # pose_shape_latent = self.pose_shape_embedding([conditions[0], conditions[1]])
         garment_latent = self.garment_embedding(conditions[2])
-        x = tf.concat([x, pose_shape_latent, garment_latent], axis=1)
+        # x = tf.concat([x, pose_shape_latent, garment_latent], axis=1)
+        x = tf.concat([x, garment_latent], axis=1)
 
         # normalize inputs
         x = self.normalize(x)

train.py

@@ -17,12 +17,13 @@ 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"] = "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)
+os.environ["CUDA_VISIBLE_DEVICES"] = "1"
+# run = wandb.init(config=tf.compat.v1.flags.FLAGS, sync_tensorboard=True, project='stylegan', name='test_conditioned_10(test-13)')
+# artifact = wandb.Artifact('code', type='code')
+# artifact.add_file('train.py')
+# artifact.add_file('losses.py')
+# artifact.add_file('stylegan2/layers/mapping_block.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)
@@ -124,13 +125,12 @@ 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)
 
-
             # scale loss
             d_loss = tf.reduce_sum(d_loss) * self.global_batch_scaler
 
@@ -141,7 +141,7 @@ class Trainer(object):
         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
@@ -162,7 +162,7 @@ class Trainer(object):
 
 
     def g_train_step(self, dist_inputs):
-        real_images, conditions, _, _, _, _ = dist_inputs[0]
+        real_images, conditions, _, _, _ = dist_inputs[0]
 
         with tf.GradientTape() as g_tape:
             # compute losses
@@ -176,7 +176,7 @@ class Trainer(object):
         return g_loss
 
     def g_train_step_reg(self, dist_inputs):
-        real_images, conditions, _, _, _, _ = dist_inputs[0]
+        real_images, conditions, _, _, _ = dist_inputs[0]
 
         with tf.GradientTape() as g_tape:
             # compute losses
@@ -198,16 +198,16 @@ class Trainer(object):
 
 
     def g_train_step_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]
 
         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)
-            c_loss = 0.0 # collision_penalty(mesh_pred, mesh_body_posed, mesh_mask)
-            m_loss = uv_mask_loss(fake_images, uv_mask)
+            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)
 
             # scale loss
             g_loss = tf.reduce_sum(g_loss) * self.global_batch_scaler
@@ -222,7 +222,8 @@ class Trainer(object):
         return g_loss, e_loss, c_loss, m_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, uv_mask = dist_inputs[0]
+        real_images, conditions, mesh_static, mesh_body_posed, mesh_mask = dist_inputs[0]
 
         with tf.GradientTape() as g_tape:
             # compute losses
@@ -233,9 +234,9 @@ class Trainer(object):
 
             mesh_pred = uv2mesh(fake_images)
 
-            e_loss = 0.1 * edge_loss(mesh_pred, mesh_static, mesh_mask)
-            c_loss = 0.0 # collision_penalty(mesh_pred, mesh_body_posed, mesh_mask)
-            m_loss = uv_mask_loss(fake_images, uv_mask)
+            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)
 
             # scale loss
             pl_penalty = tf.reduce_sum(pl_penalty) * self.global_batch_scaler
@@ -570,7 +571,7 @@ def main():
             'train_res': args['train_res'],
         }
 
-        trainer = Trainer(training_parameters, name=f'stylegan2-test-10-{args["train_res"]}x{args["train_res"]}')
+        trainer = Trainer(training_parameters, name=f'stylegan2-test-13-{args["train_res"]}x{args["train_res"]}')
         trainer.train(dist_dataset, strategy)
     return