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

working pyramid update

Browse Source
This commit is contained in:
luni
2022-09-20 22:56:44 +00:00
parent 456a8b2c37
commit 494ceaa276
5 changed files with 557 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cvae_pyramid.py
+18 -13
View File
@@ -1,5 +1,5 @@
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
os.environ["CUDA_VISIBLE_DEVICES"] = "2"
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_2_d128'
name='cvae_pyr_3_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. \
@@ -49,11 +49,13 @@ fixed gender issue and retrained body encoder\
remove shape from unposed garment\
training on 100 garments with offsets over the average of these\
l2 loss on both 128 and full res, bottom predicts base, high res predicts offsets over this\
low and high res both use same full res condition encoders, but low res also gets low res versions for main encoder input')
low and high res both use same full res condition encoders, but low res also gets low res versions for main encoder input\
not using filled in versions, but using custom MLP upscaling')
artifact = wandb.Artifact('code', type='code')
artifact.add_file('cvae_pyramid.py')
artifact.add_file('models/cvae_model.py')
artifact.add_file('models/cvae_pyramid_model.py')
artifact.add_file('models/upscaler.py')
artifact.add_file('dataset.py')
artifact.add_file('losses.py')
artifact.add_file('convnext/models/convnext_tf.py')
@@ -135,8 +137,8 @@ def train_step(batch):
uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
uv_curvature_low = tf.one_hot(batch['uv_curvature' + resolution_modifier], 7)
pred_full, mean_full, logvar_full, pred_low, mean_low, logvar_low, = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature,
batch['uv_unposed' + resolution_modifier], batch['uv_static' + resolution_modifier], batch['uv_body_posed' + resolution_modifier], uv_curvature_low])
pred_full, mean_full, logvar_full, pred_low, mean_low, logvar_low, = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask'],
batch['uv_unposed' + resolution_modifier], batch['uv_static' + resolution_modifier], batch['uv_body_posed' + resolution_modifier], uv_curvature_low, batch['uv_mask' + resolution_modifier],])
mesh_pred_full = uv2mesh(pred_full)
mesh_pred_posed_full = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_full)
@@ -152,11 +154,12 @@ 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) * batch['uv_mask'], 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) * batch['uv_mask'], axis=[1, 2])
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)
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])
@@ -236,8 +239,8 @@ 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' + resolution_modifier], 7)
pred_full, mean_full, logvar_full, pred_low, mean_low, logvar_low, = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature,
batch['uv_unposed' + resolution_modifier], batch['uv_static' + resolution_modifier], batch['uv_body_posed' + resolution_modifier], uv_curvature_low])
pred_full, mean_full, logvar_full, pred_low, mean_low, logvar_low, = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask'],
batch['uv_unposed' + resolution_modifier], batch['uv_static' + resolution_modifier], batch['uv_body_posed' + resolution_modifier], uv_curvature_low, batch['uv_mask' + resolution_modifier]])
mesh_pred_full = uv2mesh(pred_full)
mesh_pred_low = uv2mesh(pred_low)
@@ -409,7 +412,7 @@ if __name__ == '__main__':
args = vars(parser.parse_args())
batch_size = 4
low_resolution = 128
low_resolution = 32
resolution_modifier = '_' + str(low_resolution)
smpl = SMPL()
@@ -433,11 +436,12 @@ if __name__ == '__main__':
# '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,),
}
dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100', batch_size=batch_size, signature=signature_dynamic, use_filled=True, resolutions=[512, low_resolution])
static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', signature=signature_static, use_filled=True, resolutions=[512, low_resolution])
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])
static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', signature=signature_static, use_filled=False, resolutions=[512, low_resolution])
dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/valid', batch_size=batch_size, signature=signature_dynamic)
static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/valid', signature=signature_static)
@@ -447,7 +451,8 @@ if __name__ == '__main__':
freeze_encoders=True,
static_model='checkpoints/ae_static_4',
body_model='checpoints/ae_body_7',
curvature_model='checpoints/ae_curvature_7')
curvature_model='checpoints/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)
optimizer = tf.optimizers.Adam(learning_rate=0.0001)