diff --git a/.gitignore b/.gitignore index 4d09f44..469e186 100644 --- a/.gitignore +++ b/.gitignore @@ -135,5 +135,9 @@ dmypy.json .vscode/ wandb/ -output/ -output.zip \ No newline at end of file +output/* +*.zip +*.out +models/ +models_archive/ +temp.jpg diff --git a/dataset.py b/dataset.py index 58ac10b..55a635e 100644 --- a/dataset.py +++ b/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, ) ) diff --git a/inference_custom.py b/inference_custom.py index ccfa0f0..b5203a7 100644 --- a/inference_custom.py +++ b/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' }, ] diff --git a/load_models.py b/load_models.py index 6537c29..e55b0dd 100644 --- a/load_models.py +++ b/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) diff --git a/losses.py b/losses.py index ee5bb63..d9245eb 100644 --- a/losses.py +++ b/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) diff --git a/models/.gitignore b/models/.gitignore deleted file mode 100644 index cfeac87..0000000 --- a/models/.gitignore +++ /dev/null @@ -1,5 +0,0 @@ -# ignore all -* - -# except this file -!.gitignore \ No newline at end of file diff --git a/reconstruct.py b/reconstruct.py index 15811c3..14d2606 100644 --- a/reconstruct.py +++ b/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 diff --git a/stylegan2/generator.py b/stylegan2/generator.py index bd28579..efa4c5b 100644 --- a/stylegan2/generator.py +++ b/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 diff --git a/stylegan2/layers/mapping_block.py b/stylegan2/layers/mapping_block.py index 75eddc0..a86311e 100644 --- a/stylegan2/layers/mapping_block.py +++ b/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) diff --git a/temp.jpg b/temp.jpg index f8c4515..25487f1 100644 Binary files a/temp.jpg and b/temp.jpg differ diff --git a/train.py b/train.py index 76ccefb..296b8bb 100644 --- a/train.py +++ b/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