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

added skinning

Browse Source
This commit is contained in:
luni
2022-06-22 23:02:31 +00:00
parent cf4b67c5a7
commit 0e88eb341c
13 changed files with 765 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files
train.py
+52 -38
View File
@@ -15,10 +15,10 @@ from dataset import DatasetSimple
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
from custom_utils.smpl import SMPL
os.environ["CUDA_VISIBLE_DEVICES"] = "4"
run = wandb.init(config=tf.compat.v1.flags.FLAGS, sync_tensorboard=True, project='stylegan', name='test_conditioned_23', notes='learning offsets, all losses from the start, smaller generator (from 128), slower learning rate for generator, add a fading out l1 loss with gt, add tanh activation for last layer, higher weights for the extra losses, same learning rate again. added extra metrics, trained from test-22-10000')
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
run = wandb.init(config=tf.compat.v1.flags.FLAGS, sync_tensorboard=True, project='stylegan', name='run_28', notes='learning offsets, all losses from the start, smaller generator (from 128), slower learning rate for generator, add a fading out l1 loss with gt, add tanh activation for last layer, higher weights for the extra losses, same learning rate again. added extra metrics, lower discr. features too. use reduce_sum for losses divided by sum of mask. learn unposed version then do skinning. lowered collsion coeff')
artifact = wandb.Artifact('code', type='code')
artifact.add_file('train.py')
artifact.add_file('losses.py')
@@ -36,6 +36,7 @@ def initiate_models(g_params, d_params, use_custom_cuda):
uv2mesh = UV2Mesh(res=256)
smpl = SMPL()
class Trainer(object):
def __init__(self, t_params, name):
@@ -50,9 +51,9 @@ class Trainer(object):
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 = 5000
self.image_summary_step = 500
self.print_step = 100
self.save_step = 10000
self.image_summary_step = 1000
self.unsupervised_coeff = 1
self.gt_loss_coeff = 1
self.reached_max_steps = False
@@ -128,7 +129,8 @@ class Trainer(object):
return params_copy
def d_train_step(self, dist_inputs):
real_images, conditions, _, _, _, _ = dist_inputs[0]
uv_posed, conditions, _, _, _, _, _, uv_unposed = dist_inputs[0]
real_images = uv_unposed
with tf.GradientTape() as d_tape:
# compute losses
@@ -144,7 +146,8 @@ class Trainer(object):
return d_loss
def d_train_step_reg(self, dist_inputs):
real_images, conditions, _, _, _, _ = dist_inputs[0]
uv_posed, conditions, _, _, _, _, _, uv_unposed = dist_inputs[0]
real_images = uv_unposed
with tf.GradientTape() as d_tape:
# compute losses
@@ -165,7 +168,8 @@ class Trainer(object):
def g_train_step(self, dist_inputs):
real_images, conditions, _, _, _, _ = dist_inputs[0]
uv_posed, conditions, _, _, _, _, _, uv_unposed = dist_inputs[0]
real_images = uv_unposed
with tf.GradientTape() as g_tape:
# compute losses
@@ -179,7 +183,8 @@ class Trainer(object):
return g_loss
def g_train_step_reg(self, dist_inputs):
real_images, conditions, _, _, _, _ = dist_inputs[0]
uv_posed, conditions, _, _, _, _, _, uv_unposed= dist_inputs[0]
real_images = uv_unposed
with tf.GradientTape() as g_tape:
# compute losses
@@ -201,25 +206,28 @@ class Trainer(object):
def g_train_step_extra(self, dist_inputs):
real_images, conditions, mesh_static, mesh_body_posed, mesh_mask, mesh_posed = dist_inputs[0]
uv_posed, conditions, mesh_static, mesh_body_posed, mesh_mask, mesh_posed, mesh_unposed, uv_unposed = dist_inputs[0]
real_images = uv_unposed
uv_mask = conditions[3]
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 = 1e4 * edge_loss(mesh_pred, mesh_static, mesh_mask)
c_loss = 1e5 * collision_penalty(mesh_pred, mesh_body_posed, mesh_mask)
mesh_pred_unposed = uv2mesh(fake_images)
mesh_pred = smpl.run_garment(conditions[0], conditions[1], mesh_pred_unposed)
e_loss = 1e4 * edge_loss(mesh_pred_unposed, mesh_static, mesh_mask)
c_loss = 1e2 * collision_penalty(mesh_pred, mesh_body_posed, mesh_mask)
l1_loss = 1e2 * tf.reduce_mean(tf.reduce_sum((tf.math.abs(real_images - fake_images)), axis=1) * uv_mask)
euclid_error = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(mesh_mask[:,:,tf.newaxis] * (mesh_posed - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(mesh_mask)
ms_ssim = tf.image.ssim_multiscale(tf.transpose(real_images, [0,2,3,1]) + 1, tf.transpose(fake_images, [0,2,3,1]) + 1, max_val=2)
# 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
l1_loss = tf.reduce_sum(l1_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
# l1_loss = tf.reduce_sum(l1_loss) * self.global_batch_scaler
ms_ssim = tf.reduce_sum(ms_ssim) * self.global_batch_scaler
g_loss = g_loss + self.unsupervised_coeff * (e_loss + c_loss) + self.gt_loss_coeff * l1_loss
@@ -230,7 +238,8 @@ class Trainer(object):
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, mesh_posed = dist_inputs[0]
uv_posed, conditions, mesh_static, mesh_body_posed, mesh_mask, mesh_posed, mesh_unposed, uv_unposed = dist_inputs[0]
real_images = uv_unposed
uv_mask = conditions[3]
with tf.GradientTape() as g_tape:
@@ -240,20 +249,21 @@ class Trainer(object):
self.pl_minibatch_shrink, self.pl_denorm, self.pl_decay, conditions=conditions)
pl_penalty = pl_penalty * self.pl_weight * self.g_opt['reg_interval']
mesh_pred = uv2mesh(fake_images)
mesh_pred_unposed = uv2mesh(fake_images)
mesh_pred = smpl.run_garment(conditions[0], conditions[1], mesh_pred_unposed)
e_loss = 1e4 * edge_loss(mesh_pred, mesh_static, mesh_mask)
c_loss = 1e5 * collision_penalty(mesh_pred, mesh_body_posed, mesh_mask)
l1_loss = 1e2 * tf.reduce_mean(tf.reduce_sum((tf.math.abs(real_images - fake_images)), axis=1) * uv_mask)
e_loss = 1e4 * edge_loss(mesh_pred_unposed, mesh_static, mesh_mask)
c_loss = 1e2 * collision_penalty(mesh_pred, mesh_body_posed, mesh_mask)
l1_loss = 1e2 * tf.reduce_sum(tf.reduce_sum((tf.math.abs(real_images - fake_images)), axis=1) * uv_mask) / tf.reduce_sum(uv_mask)
euclid_error = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(mesh_mask[:,:,tf.newaxis] * (mesh_posed - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(mesh_mask)
ms_ssim = tf.image.ssim_multiscale(tf.transpose(real_images, [0,2,3,1]) + 1, tf.transpose(fake_images, [0,2,3,1]) + 1, max_val=2)
# 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
l1_loss = tf.reduce_sum(l1_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
# l1_loss = tf.reduce_sum(l1_loss) * self.global_batch_scaler
ms_ssim = tf.reduce_sum(ms_ssim) * self.global_batch_scaler
# combine
@@ -296,9 +306,9 @@ class Trainer(object):
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_l1_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[3], axis=None)
mean_e_loss = strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_losses[1], axis=None)
mean_c_loss = strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_losses[2], axis=None)
mean_l1_loss = strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_losses[3], axis=None)
mean_euclid_error = strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_losses[4], axis=None)
mean_ms_ssim = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[5], axis=None)
return mean_g_loss, mean_e_loss, mean_c_loss, mean_l1_loss, mean_euclid_error, mean_ms_ssim
@@ -309,9 +319,9 @@ 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)
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)
mean_l1_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[5], axis=None)
mean_e_loss = strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_losses[3], axis=None)
mean_c_loss = strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_losses[4], axis=None)
mean_l1_loss = strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_losses[5], axis=None)
mean_euclid_error = strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_losses[6], axis=None)
mean_ms_ssim = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[7], axis=None)
return mean_g_loss, mean_g_gan_loss, mean_pl_penalty, mean_e_loss, mean_c_loss, mean_l1_loss, mean_euclid_error, mean_ms_ssim
@@ -440,7 +450,8 @@ 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):])
mesh_unposed = uv2mesh(summary_image[:, :, :self.train_res, int(self.train_res/2):])
mesh = smpl.run_garment(batch[1][0], batch[1][1], mesh_unposed)
# convert to tensor image
summary_image = self.convert_per_replica_image(summary_image, strategy)
@@ -449,7 +460,7 @@ class Trainer(object):
template = uv2mesh(batch[1][2])
# gt = uv2mesh(batch[0])
plot_meshes(mesh_writer, garment=mesh, mesh_gt=batch[5], template=template, body=batch[3], mask=batch[4].numpy(), step=step)
plot_meshes(mesh_writer, garment_unposed=mesh_unposed, garment=mesh, mesh_gt=batch[6], 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()
@@ -503,14 +514,16 @@ def filter_resolutions_featuremaps(resolutions, featuremaps, res):
return filtered_resolutions, filtered_featuremaps
def plot_meshes(writer, garment, mesh_gt, template, body, mask, step, nr_offset=0):
def plot_meshes(writer, garment_unposed, garment, mesh_gt, template, body, mask, step, nr_offset=0):
body = body.numpy()
template = template.numpy()
mesh_gt = mesh_gt.numpy()
garment_unposed = garment_unposed.numpy()
F = np.load("custom_utils/smpl_faces.npy")
garment, faces = maskFacesBatch(garment.numpy(), F, mask.astype(bool))
template = [template[i, mask.astype(bool)[i,:], :] for i in range(len(garment))]
mesh_gt = [mesh_gt[i, mask.astype(bool)[i,:], :] for i in range(len(garment))]
garment_unposed = [garment_unposed[i, mask.astype(bool)[i,:], :] for i in range(len(garment))]
collision = [np.concatenate((body[i], garment[i]), axis=0) for i in range(len(garment))]
collision_faces = [np.concatenate((F, faces[i] + 14475), axis=0) for i in range(len(garment))]
@@ -524,7 +537,7 @@ def plot_meshes(writer, garment, mesh_gt, template, body, mask, step, nr_offset=
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,:,:])
for i in range(len(garment)):
plotly_fig = plot_3d(mesh_gt[i], garment[i], faces[i], collision[i], collision_faces[i], collision_color[i])
plotly_fig = plot_3d(mesh_gt[i], garment_unposed[i], faces[i], collision[i], collision_faces[i], collision_color[i])
wandb.log({"garment_" + str(i): plotly_fig})
def main():
@@ -553,7 +566,8 @@ def main():
# network params
resolutions = [4, 8, 16, 32, 64, 128, 256, 512, 1024]
g_featuremaps = [128, 128, 128, 128, 128, 64, 32, 16, 8]
d_featuremaps = [512, 512, 512, 512, 512, 256, 128, 64, 32]
d_featuremaps = [128, 128, 128, 128, 128, 64, 32, 16, 8]
# d_featuremaps = [512, 512, 512, 512, 512, 256, 128, 64, 32]
# featuremaps = [512, 512, 512, 512, 512, 256, 128, 64, 32]
# featuremaps = [64, 64, 64, 64, 64, 32, 16, 8, 4]
train_resolutions, train_g_featuremaps = filter_resolutions_featuremaps(resolutions, g_featuremaps, args['train_res'])
@@ -596,7 +610,7 @@ def main():
'use_custom_cuda': args['use_custom_cuda'],
'model_base_dir': args['model_base_dir'],
'enable_unsupervised_losses': 0,
'disable_l1_loss': 10000,
'disable_l1_loss': 30000,
# network params
'g_params': g_params,
@@ -611,7 +625,7 @@ def main():
'train_res': args['train_res'],
}
trainer = Trainer(training_parameters, name=f'stylegan2-test-23-{args["train_res"]}x{args["train_res"]}')
trainer = Trainer(training_parameters, name=f'stylegan2-test-28-{args["train_res"]}x{args["train_res"]}')
trainer.train(dist_dataset, strategy)
return