diff --git a/cvae_pyramid.py b/cvae_pyramid.py index 28f792e..fc9dada 100644 --- a/cvae_pyramid.py +++ b/cvae_pyramid.py @@ -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_1_d128' +name='cvae_pyr_2_d128' 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. \ @@ -48,7 +48,8 @@ use mean l2 loss. small batch. l2 * 1000, balancing to 1 from 0.01\ 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') +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') artifact = wandb.Artifact('code', type='code') artifact.add_file('cvae_pyramid.py') artifact.add_file('models/cvae_model.py') @@ -250,19 +251,21 @@ def plot_everything(batch: dict, type: str): img = tf.concat((uv_static_img, uv_body_posed_img, orig_img, pred_img, pred_low_up), axis=2) imgs = [wandb.Image((img[i] + 1) / 2.0) 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") mask = batch['mesh_mask'].numpy().astype(bool) garment_unposed, faces = maskFacesBatch(mesh_pred_full.numpy(), F, mask) 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])] body = batch['mesh_body_posed'].numpy() mesh_posed_full = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_full).numpy() - 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])]) - garment_posed = [mesh_posed_full[i, mask[i], :] + shape_offsets.numpy()[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))] @@ -439,7 +442,12 @@ if __name__ == '__main__': static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/valid', signature=signature_static) - model = Pyramid(latent_dim=128, low_res=low_resolution) + 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') # 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) diff --git a/models/cvae_model.py b/models/cvae_model.py index d9241bb..df8f3eb 100644 --- a/models/cvae_model.py +++ b/models/cvae_model.py @@ -258,21 +258,31 @@ class CVAE32(tf.keras.Model): static_model='', body_model='', curvature_model='', - freeze_encoders=False): + freeze_encoders=False, + use_full_res_encoders=True): super(CVAE32, self).__init__() self.latent_dim = latent_dim - model : StaticAE32 = StaticAE32() + if use_full_res_encoders: + model : StaticAE = StaticAE() + else: + model : StaticAE32 = StaticAE32() if static_model != '': load_model(model, static_model) self.static_enc = tf.keras.Sequential([model.encoder, model.bottleneck]) - model : BodyAE32 = BodyAE32() + if use_full_res_encoders: + model : BodyAE = BodyAE() + else: + model : BodyAE32 = BodyAE32() if body_model != '': load_model(model, body_model) self.body_enc = tf.keras.Sequential([model.encoder, model.bottleneck]) - model : CurvatureAE32 = CurvatureAE32() + if use_full_res_encoders: + model : CurvatureAE = CurvatureAE() + else: + model : CurvatureAE32 = CurvatureAE32() if curvature_model != '': load_model(model, curvature_model) self.curv_enc = tf.keras.Sequential([model.encoder, model.bottleneck]) @@ -429,21 +439,31 @@ class CVAE128(tf.keras.Model): static_model='', body_model='', curvature_model='', - freeze_encoders=False): + freeze_encoders=False, + use_full_res_encoders=True): super(CVAE128, self).__init__() self.latent_dim = latent_dim - model : StaticAE128 = StaticAE128() + if use_full_res_encoders: + model : StaticAE = StaticAE() + else: + model : StaticAE128 = StaticAE128() if static_model != '': load_model(model, static_model) self.static_enc = tf.keras.Sequential([model.encoder, model.bottleneck]) - - model : BodyAE128 = BodyAE128() + + if use_full_res_encoders: + model : BodyAE = BodyAE() + else: + model : BodyAE128 = BodyAE128() if body_model != '': load_model(model, body_model) self.body_enc = tf.keras.Sequential([model.encoder, model.bottleneck]) - model : CurvatureAE128 = CurvatureAE128() + if use_full_res_encoders: + model : CurvatureAE = CurvatureAE() + else: + model : CurvatureAE128 = CurvatureAE128() if curvature_model != '': load_model(model, curvature_model) self.curv_enc = tf.keras.Sequential([model.encoder, model.bottleneck]) @@ -515,11 +535,11 @@ class CVAE128(tf.keras.Model): def call(self, inputs, training=None, zero_latent=False): - uv_unposed, uv_template, uv_body_posed, uv_curv = inputs + uv_unposed, uv_template, uv_body_posed, uv_curv, uv_template_high, uv_body_posed_high, uv_curv_high = inputs uv_unposed, uv_template, uv_body_posed = tf.transpose(uv_unposed, [0,2,3,1]), tf.transpose(uv_template, [0,2,3,1]), tf.transpose(uv_body_posed, [0,2,3,1]) - z_static = self.static_enc(uv_template) - z_body = self.body_enc(uv_body_posed) - z_curv = self.curv_enc(uv_curv) + z_static = self.static_enc(uv_template_high) + z_body = self.body_enc(uv_body_posed_high) + z_curv = self.curv_enc(uv_curv_high) x = tf.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], axis=-1) diff --git a/models/cvae_pyramid_model.py b/models/cvae_pyramid_model.py index dfc1fd7..ff12319 100644 --- a/models/cvae_pyramid_model.py +++ b/models/cvae_pyramid_model.py @@ -12,15 +12,39 @@ from models.cvae_model import CVAE, CVAE128, CVAE32 class Pyramid(tf.keras.Model): - def __init__(self, low_res=32, latent_dim=128, avg_garment_path=''): + def __init__(self, low_res=32, + latent_dim=128, + avg_garment_path='', + static_model='', + body_model='', + curvature_model='', + freeze_encoders=False, + use_full_res_encoders=True): super(Pyramid, self).__init__() if low_res == 32: - self.model_low = CVAE32(latent_dim=latent_dim) + self.model_low = CVAE32(latent_dim=latent_dim, + static_model=static_model, + body_model=body_model, + curvature_model=curvature_model, + freeze_encoders=freeze_encoders, + use_full_res_encoders=use_full_res_encoders) if low_res == 128: - self.model_low = CVAE128(latent_dim=latent_dim) - self.model_full = CVAE(latent_dim=latent_dim) + self.model_low = CVAE128(latent_dim=latent_dim, + static_model=static_model, + body_model=body_model, + curvature_model=curvature_model, + freeze_encoders=freeze_encoders, + use_full_res_encoders=use_full_res_encoders) + + self.model_full = CVAE(latent_dim=latent_dim, + static_model=static_model, + body_model=body_model, + curvature_model=curvature_model, + freeze_encoders=freeze_encoders) + self.use_full_res_encoders = use_full_res_encoders + if avg_garment_path != '': self.avg_garment = np.load(avg_garment_path) else: @@ -29,7 +53,10 @@ class Pyramid(tf.keras.Model): def call(self, inputs, training=None): uv_unposed, uv_static, uv_body_posed, uv_curv, uv_unposed_low, uv_static_low, uv_body_posed_low, uv_curv_low = inputs - pred_low, mean_low, logvar_low, _ = self.model_low([uv_unposed_low, uv_static_low, uv_body_posed_low, uv_curv_low], training=training) + if self.use_full_res_encoders: + pred_low, mean_low, logvar_low, _ = self.model_low([uv_unposed_low, uv_static_low, uv_body_posed_low, uv_curv_low, uv_static, uv_body_posed, uv_curv], training=training) + else: + pred_low, mean_low, logvar_low, _ = self.model_low([uv_unposed_low, uv_static_low, uv_body_posed_low, uv_curv_low], training=training) x = tf.image.resize(pred_low, (512, 256), tf.image.ResizeMethod.LANCZOS5)