luni/cvpr
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

plotting not working, losses fixed?

Browse Source
This commit is contained in:
luni
2022-06-10 15:32:25 +00:00
parent 31607cc754
commit 94e693d7f6
5 changed files with 108 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+3 -1
View File
@@ -134,4 +134,6 @@ dmypy.json
# vscode
.vscode/
wandb/
wandb/
output/
output.zip
custom_utils/custom_utils.py
+23 -1
View File
@@ -62,4 +62,26 @@ 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
+18 -9
View File
@@ -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',
'use_custom_cuda': False,
'folder': '/workspace/Checkpoints/checkpoint_52000/output',
'ckpt_dir': 'models/stylegan2-test-5-512x512',
'use_custom_cuda': True,
'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
+8 -9
View File
@@ -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]))
loss_1 = collision_loss(V, B, thr=.004, stop_gradient=False)
loss_2 = collision_penalty(V, B, B_vn)
m = tf.ones((4, 14475))
loss_1 = collision_loss(V, B, m, thr=.004)
print()
train.py
+56 -34
View File
@@ -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"
wandb.init(config=tf.compat.v1.flags.FLAGS, sync_tensorboard=True, project='stylegan', name='test_conditioned_1')
os.environ["CUDA_VISIBLE_DEVICES"] = "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)
c_loss = collision_penalty(mesh_pred, mesh_body_posed)
# 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
# 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# + 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)
c_loss = collision_penalty(mesh_pred, mesh_body_posed)
# 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
# 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# + 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_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)
return mean_d_loss, mean_e_loss, mean_c_loss
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_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,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_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
# 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_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