large refactor
This commit is contained in:
+28
-26
@@ -1,5 +1,5 @@
|
||||
import os
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = "2"
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
|
||||
import tensorflow as tf
|
||||
|
||||
gpus = tf.config.experimental.list_physical_devices('GPU')
|
||||
@@ -37,7 +37,7 @@ from custom_utils.smpl import SMPL
|
||||
from models.cvae_pyramid_model import Pyramid
|
||||
|
||||
|
||||
name='cvae_pyr_3_d32'
|
||||
name='cvae_pyr_4_d32'
|
||||
def init_wandb():
|
||||
run = wandb.init(project='vae', name=name,
|
||||
notes='learn offsets. smaller flattened latent vectors. with 25-85k pretraining. smaller batchsize. fixed balancing of reconst and kl loss for baseline. \
|
||||
@@ -130,7 +130,7 @@ def kl_loss(mean, logvar):
|
||||
return kl_loss
|
||||
|
||||
|
||||
@tf.function
|
||||
# @tf.function
|
||||
def train_step(batch):
|
||||
|
||||
with tf.GradientTape() as tape:
|
||||
@@ -154,17 +154,17 @@ def train_step(batch):
|
||||
|
||||
e_loss = tf.reduce_mean(edge_loss(mesh_pred_full, mesh_orig, batch['mesh_mask']))
|
||||
c_loss = tf.reduce_mean(collision_penalty(mesh_pred_posed_full, batch['mesh_body_posed'], batch['mesh_mask']))
|
||||
l1_loss = tf.reduce_sum(tf.reduce_sum((tf.math.abs(tf.transpose(batch['uv_unposed'], [0,2,3,1]) - pred_full)), axis=-1), axis=[1, 2])
|
||||
l2_loss = 1e3 * tf.reduce_mean(tf.reduce_sum(((tf.transpose(batch['uv_unposed'], [0,2,3,1]) - pred_full) ** 2), axis=-1), axis=[1, 2])
|
||||
ms_ssim = tf.image.ssim_multiscale(pred_full + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=3)
|
||||
l1_loss = tf.reduce_sum(tf.reduce_sum((tf.math.abs(batch['uv_unposed'] - pred_full)), axis=-1), axis=[1, 2])
|
||||
l2_loss = 1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed'] - pred_full) ** 2), axis=-1), axis=[1, 2])
|
||||
ms_ssim = tf.image.ssim_multiscale(pred_full + 1, batch['uv_unposed'] + 1, max_val=3)
|
||||
kl = kl_loss(mean_full, logvar_full)
|
||||
|
||||
pred_low = pred_low * batch['uv_mask' + resolution_modifier][:,:,:,tf.newaxis]
|
||||
e_loss_low = tf.reduce_mean(edge_loss(mesh_pred_low, mesh_orig, batch['mesh_mask']))
|
||||
c_loss_low = tf.reduce_mean(collision_penalty(mesh_pred_posed_low, batch['mesh_body_posed'], batch['mesh_mask']))
|
||||
l1_loss_low = tf.reduce_sum(tf.reduce_sum((tf.math.abs(tf.transpose(batch['uv_unposed' + resolution_modifier], [0,2,3,1]) - pred_low)), axis=-1), axis=[1, 2])
|
||||
l2_loss_low = 1e3 * tf.reduce_mean(tf.reduce_sum(((tf.transpose(batch['uv_unposed' + resolution_modifier], [0,2,3,1]) - pred_low) ** 2), axis=-1), axis=[1, 2])
|
||||
ms_ssim_low = tf.image.ssim(pred_low + 1, tf.transpose(batch['uv_unposed' + resolution_modifier], [0,2,3,1]) + 1, max_val=3)
|
||||
l1_loss_low = tf.reduce_sum(tf.reduce_sum((tf.math.abs(batch['uv_unposed' + resolution_modifier] - pred_low)), axis=-1), axis=[1, 2])
|
||||
l2_loss_low = 1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed' + resolution_modifier] - pred_low) ** 2), axis=-1), axis=[1, 2])
|
||||
ms_ssim_low = tf.image.ssim(pred_low + 1, batch['uv_unposed' + resolution_modifier] + 1, max_val=3)
|
||||
kl_low = kl_loss(mean_full, logvar_full)
|
||||
|
||||
reconstr_loss = l2_loss + l2_loss_low
|
||||
@@ -246,9 +246,9 @@ def plot_everything(batch: dict, type: str):
|
||||
mesh_pred_low = uv2mesh(pred_low)
|
||||
mesh_orig = batch['mesh_unposed']
|
||||
|
||||
orig_img = tf.transpose(batch['uv_unposed'], [0,2,3,1])
|
||||
uv_static_img = tf.transpose(batch['uv_static'], [0,2,3,1])
|
||||
uv_body_posed_img = tf.transpose(batch['uv_body_posed'], [0,2,3,1])
|
||||
orig_img = batch['uv_unposed']
|
||||
uv_static_img = batch['uv_static']
|
||||
uv_body_posed_img = batch['uv_body_posed']
|
||||
pred_img = pred_full
|
||||
pred_low_up = tf.image.resize(pred_low, (512, 256), tf.image.ResizeMethod.LANCZOS5)
|
||||
img = tf.concat((uv_static_img, uv_body_posed_img, orig_img, pred_img, pred_low_up), axis=2)
|
||||
@@ -264,12 +264,14 @@ def plot_everything(batch: dict, type: str):
|
||||
garment_unposed_low, faces = maskFacesBatch(mesh_pred_low.numpy(), F, mask)
|
||||
garment_unposed = [garment_unposed[i] + shape_offsets.numpy()[i, mask[i], :] for i in range(mask.shape[0])]
|
||||
garment_unposed_low = [garment_unposed_low[i] + shape_offsets.numpy()[i, mask[i], :] for i in range(mask.shape[0])]
|
||||
gt = [mesh_orig[i].numpy()[mask[i], :] for i in range(pred_full.shape[0])]
|
||||
gt = [mesh_orig[i].numpy()[mask[i], :] + shape_offsets.numpy()[i, mask[i], :] for i in range(mask.shape[0])]
|
||||
|
||||
for i in range(mask.shape[0]):
|
||||
mesh_pred_full[i, mask[i], :] += shape_offsets.numpy()[i, mask[i], :]
|
||||
body = batch['mesh_body_posed'].numpy()
|
||||
mesh_posed_full = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_full).numpy()
|
||||
|
||||
garment_posed = [mesh_posed_full[i, mask[i], :] + shape_offsets.numpy()[i, mask[i], :] for i in range(body.shape[0])]
|
||||
garment_posed = [mesh_posed_full[i, mask[i], :] for i in range(body.shape[0])]
|
||||
collision = [np.concatenate((body[i], garment_posed[i]), axis=0) for i in range(len(garment_posed))]
|
||||
collision_faces = [np.concatenate((F, faces[i] + 14475), axis=0) for i in range(len(garment_posed))]
|
||||
collision_color = [np.concatenate((np.ones_like(body[i], dtype=np.int32)*100, np.ones_like(garment_posed[i], dtype=np.int32)*230), axis=0) for i in range(len(garment_posed))]
|
||||
@@ -405,7 +407,7 @@ def sample(folder):
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
init_wandb()
|
||||
# init_wandb()
|
||||
allow_memory_growth()
|
||||
parser = argparse.ArgumentParser(description='')
|
||||
parser.add_argument('--name', type=str)
|
||||
@@ -420,24 +422,24 @@ if __name__ == '__main__':
|
||||
factory = DatasetFactory()
|
||||
signature_dynamic = { 'id': (),
|
||||
'pose_shape': (72 + 11,),
|
||||
# 'uv_posed': (3, 512, 256,),
|
||||
'uv_body_posed': (3, 512, 256,),
|
||||
'uv_body_posed' + resolution_modifier: (3, low_resolution, low_resolution//2,),
|
||||
# 'uv_posed': (512, 256, 3),
|
||||
'uv_body_posed': (512, 256, 3),
|
||||
'uv_body_posed' + resolution_modifier: (low_resolution, low_resolution//2, 3),
|
||||
'mesh_body_posed': (14475, 3,),
|
||||
'mesh_posed': (14475, 3,),
|
||||
'mesh_unposed': (14475, 3,),
|
||||
'uv_unposed': (3, 512, 256,),
|
||||
'uv_unposed' + resolution_modifier: (3, low_resolution, low_resolution//2,),
|
||||
'uv_unposed': (512, 256, 3),
|
||||
'uv_unposed' + resolution_modifier: (low_resolution, low_resolution//2, 3),
|
||||
'uv_curvature': (512, 256),
|
||||
'uv_curvature' + resolution_modifier: (low_resolution, low_resolution//2),
|
||||
}
|
||||
signature_static = { 'uv_static': (3, 512, 256,),
|
||||
'uv_static' + resolution_modifier: (3, low_resolution, low_resolution//2,),
|
||||
signature_static = { 'uv_static': (512, 256, 3),
|
||||
'uv_static' + resolution_modifier: (low_resolution, low_resolution//2, 3),
|
||||
# 'mesh_static': (14475, 3,),
|
||||
'mesh_mask': (14475,),
|
||||
'uv_mask': (512, 256,),
|
||||
'uv_mask' + resolution_modifier: (low_resolution, low_resolution//2,),
|
||||
# 'uv_shaped_body': (3, 512, 256,),
|
||||
# 'uv_shaped_body': (512, 256, 3),
|
||||
}
|
||||
|
||||
dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100', batch_size=batch_size, signature=signature_dynamic, use_filled=False, resolutions=[512, low_resolution])
|
||||
@@ -449,9 +451,9 @@ if __name__ == '__main__':
|
||||
model = Pyramid(latent_dim=128,
|
||||
low_res=low_resolution,
|
||||
freeze_encoders=True,
|
||||
static_model='checkpoints/ae_static_4',
|
||||
body_model='checpoints/ae_body_7',
|
||||
curvature_model='checpoints/ae_curvature_7',
|
||||
static_model='checkpoints/ae_static_6',
|
||||
body_model='checkpoints/ae_body_8',
|
||||
curvature_model='checkpoints/ae_curvature_7',
|
||||
upscaler_model='checkpoints/up_2_3')
|
||||
|
||||
# lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(0.0001, decay_steps=30000, decay_rate=0.1, staircase=False)
|
||||
|
||||
Reference in New Issue
Block a user