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

pyramid with common encoders per level

Browse Source
This commit is contained in:
luni
2022-09-16 20:51:31 +00:00
parent eac8aad8ab
commit 37df8468e2
3 changed files with 81 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cvae_pyramid.py
+16 -8
View File
@@ -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)
models/cvae_model.py
+26 -6
View File
@@ -258,20 +258,30 @@ 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
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])
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])
if use_full_res_encoders:
model : CurvatureAE = CurvatureAE()
else:
model : CurvatureAE32 = CurvatureAE32()
if curvature_model != '':
load_model(model, curvature_model)
@@ -429,20 +439,30 @@ 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
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])
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])
if use_full_res_encoders:
model : CurvatureAE = CurvatureAE()
else:
model : CurvatureAE128 = CurvatureAE128()
if curvature_model != '':
load_model(model, curvature_model)
@@ -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)
models/cvae_pyramid_model.py
+31 -4
View File
@@ -12,14 +12,38 @@ 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)
@@ -29,6 +53,9 @@ 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
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)