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

making ready for skirts

Browse Source
This commit is contained in:
luni
2022-09-30 21:53:17 +00:00
parent 6a2f9ff4ff
commit dca5368704
22 changed files with 700 additions and 219 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cvae_low.py
+29 -30
View File
@@ -27,7 +27,12 @@ import plotly.express as ex
from PIL import Image
from tf_utils import check_tf_version, allow_memory_growth, split_gpu_for_testing
from losses import d_logistic, d_logistic_r1_reg, g_logistic_non_saturating, g_logistic_ns_pathreg, edge_loss, uv_mask_loss, bend_loss
if config['is_skirt']:
from losses import edge_loss_skirt as edge_loss
else:
from losses import edge_loss
# from losses import collision_penalty as collision_penalty
from losses import collision_loss as collision_penalty
from custom_utils.custom_utils import batch_to_dict, load_model, plot_3d
@@ -51,7 +56,7 @@ def init_wandb():
artifact.add_file('convnext/models/convnext_tf.py')
run.log_artifact(artifact)
uv2mesh = UV2Mesh(res=256)
uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'])
def log_normal_pdf(sample, mean, logvar, raxis=1):
@@ -71,10 +76,8 @@ def kl_loss(mean, logvar):
def train_step(batch):
with tf.GradientTape() as tape:
uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
uv_curvature_low = tf.one_hot(batch['uv_curvature' + res_modifier], 7)
pred, mean, logvar, z = model([batch['uv_unposed' + res_modifier], batch['uv_static' + res_modifier], batch['uv_body_posed' + res_modifier], uv_curvature_low,
batch['uv_mask' + res_modifier], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask']], use_mask=True)
pred, mean, logvar = model([batch['uv_unposed' + res_modifier], batch['uv_static' + res_modifier], batch['uv_body_posed' + res_modifier], batch['uv_curvature' + res_modifier],
batch['uv_mask' + res_modifier], batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']])
pred_enhanced = enhance(upscale_enhancer, pred, batch['uv_mask' + res_modifier], training=False)
pred_full_res = tf.image.resize(pred_enhanced, size=(512, 256), method=tf.image.ResizeMethod.LANCZOS5) * batch['uv_mask'][:,:,:,tf.newaxis]
@@ -91,22 +94,18 @@ def train_step(batch):
c_loss = collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
l1_loss = tf.reduce_sum(tf.reduce_sum((tf.math.abs(batch['uv_unposed' + res_modifier] - pred)), axis=-1), axis=[1, 2])
l2_loss = 1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed' + res_modifier] - pred) ** 2), axis=-1), axis=[1, 2])
l2_loss_full = 1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed'] - pred_full_res) ** 2), axis=-1), axis=[1, 2])
if resolution == 32:
ms_ssim = tf.image.ssim(pred + 1, batch['uv_unposed' + res_modifier] + 1, max_val=3)
ms_ssim_full = tf.image.ssim(pred_full_res + 1, batch['uv_unposed'] + 1, max_val=3)
else:
ms_ssim = tf.image.ssim_multiscale(pred + 1, batch['uv_unposed' + res_modifier] + 1, max_val=3, power_factors=(0.5, 0.3, 0.1, 0.1))
reconstr_loss = l2_loss_full #+ e_loss + c_loss
reconstr_loss = l2_loss #+ e_loss + c_loss
reconstr_loss = tf.reduce_mean(reconstr_loss)
kl = kl_loss(mean, logvar)
loss = reconstr_loss + balancing_coeff * kl
ms_ssim = tf.reduce_mean(ms_ssim)
ms_ssim_full = tf.reduce_mean(ms_ssim_full)
l1_loss = tf.reduce_mean(l1_loss)
l2_loss = tf.reduce_mean(l2_loss)
l2_loss_full = tf.reduce_mean(l2_loss_full)
e_loss = tf.reduce_mean(e_loss)
c_loss = tf.reduce_mean(c_loss)
mean = tf.reduce_mean(mean)
@@ -116,7 +115,6 @@ def train_step(batch):
gradients = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(gradients, model.trainable_variables))
return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/collision_loss': c_loss, 'train/l1_loss': l1_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim,
'train/l2_loss_full_res': l2_loss_full, 'train/ms_ssim_full_res': ms_ssim_full,
'train/reconstr_loss': reconstr_loss,
'train/mean': mean, 'train/logvar': logvar, 'train/kl_loss': kl, 'train/var': np.e ** logvar,
'train/reg_coeff': reg_coeff, 'train/l2_loss': l2_loss }
@@ -163,29 +161,28 @@ def valid_step(batch):
def plot_everything(batch: dict, type: str):
uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
uv_curvature_low = tf.one_hot(batch['uv_curvature' + res_modifier], 7)
pred, mean, logvar, z = model([batch['uv_unposed' + res_modifier], batch['uv_static' + res_modifier], batch['uv_body_posed' + res_modifier], uv_curvature_low,
batch['uv_mask' + res_modifier], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask']], use_mask=True)
pred, mean, logvar = model([batch['uv_unposed' + res_modifier], batch['uv_static' + res_modifier], batch['uv_body_posed' + res_modifier], batch['uv_curvature' + res_modifier],
batch['uv_mask' + res_modifier], batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']])
pred = enhance(upscale_enhancer, pred, batch['uv_mask' + res_modifier], training=False)
mesh_pred = uv2mesh(pred).numpy()
pred_full_res = tf.image.resize(pred, size=(512, 256), method=tf.image.ResizeMethod.LANCZOS5) * batch['uv_mask'][:,:,:,tf.newaxis]
mesh_pred = uv2mesh(pred_full_res).numpy()
mesh_orig = batch['mesh_unposed']
orig_img = batch['uv_unposed' + res_modifier]
uv_static_img = batch['uv_static' + res_modifier]
uv_body_posed_img = batch['uv_body_posed' + res_modifier]
pred_img = pred.numpy()
pred_img[pred_img > 1] = 0
pred_img[pred_img < -1] = 0
img = tf.concat((uv_static_img, uv_body_posed_img, orig_img, pred_img), axis=2)
imgs = [wandb.Image((img[i] + 1) / 2.0) for i in range(4)]
orig_img = (batch['uv_unposed'].numpy() + 1) / 2
uv_static_img = (batch['uv_static'].numpy() + 1) / 2
uv_body_posed_img = (batch['uv_body_posed'].numpy() + 1) / 2
uv_curvature = (np.tile(batch['uv_curvature'].numpy().argmax(-1)[:,:,:,np.newaxis], [1,1,1,3]) / 6)
pred_img = pred_full_res.numpy()
pred_img = (pred_img - pred_img.min(axis=(1,2,3))[:,np.newaxis,np.newaxis,np.newaxis]) / (pred_img.max(axis=(1,2,3))[:,np.newaxis,np.newaxis,np.newaxis] - pred_img.min(axis=(1,2,3))[:,np.newaxis,np.newaxis,np.newaxis])
img = tf.concat((uv_static_img, uv_body_posed_img, uv_curvature, orig_img, pred_img), axis=2)
imgs = [wandb.Image(img[i]) for i in range(4)]
shape = batch['pose_shape'][:, 72:]
gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
F = np.load("custom_utils/smpl_faces.npy")
F = np.load("custom_utils/smpl_faces" + skirt_modifier + ".npy")
mask = batch['mesh_mask'].numpy().astype(bool)
garment_unposed, faces = maskFacesBatch(mesh_pred, F, mask)
garment_unposed = [garment_unposed[i] + shape_offsets.numpy()[i, mask[i], :] for i in range(mask.shape[0])]
@@ -366,11 +363,13 @@ if __name__ == '__main__':
smpl = SMPL()
resolution = 32
resolution = config['resolution']
res_modifier = ''
if resolution != 512:
res_modifier = '_' + str(resolution)
skirt_modifier = '_skirt' if config['is_skirt'] else ''
factory = DatasetFactory()
dataset_train = factory.get_dataset(folder=config['train_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[512, config['resolution']], in_memory=config['dataset_in_memory'])
@@ -380,7 +379,7 @@ if __name__ == '__main__':
model = CVAE(config)
upscale_enhancer = get_upscale_enhancer('checkpoints/up_2_3', res=resolution)
upscale_enhancer = get_upscale_enhancer(config['upscale_enhancer'], res=resolution)
optimizer = tf.optimizers.Adam(learning_rate=config['lr'])
balancing_coeff = config['starting_balancing_coeff']