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

pro sampling avaiable, updated cvae

Browse Source
This commit is contained in:
luni
2022-09-22 16:32:37 +00:00
parent c5ae37aacc
commit a124f95f1f
4 changed files with 297 additions and 118 deletions
Download Patch File Download Diff File Expand all files Collapse all files
models/cvae_model.py
+47 -15
View File
@@ -126,7 +126,7 @@ class CVAE(tf.keras.Model):
# self.shuffle = tf.keras.layers.Lambda(lambda x: tf.stack([tf.gather(x[i], tf.random.shuffle(tf.range(x[i].shape[0]))) for i in range(x.shape[0])]))
# @tf.function
def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask, nr_of_samples=4, fix_noise=False, zero_latent=False):
def sample_single(self, uv_template, uv_body_posed, uv_curv, uv_mask, nr_of_samples=4, fix_noise=False):
if not fix_noise:
eps = tf.random.normal(shape=(nr_of_samples, self.latent_dim))
else:
@@ -147,9 +147,6 @@ class CVAE(tf.keras.Model):
z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
z_cond = self.cond_encoder(z_cond)
if zero_latent:
eps = tf.zeros_like(eps)
z = tf.concat([eps, z_cond], axis=-1)
x = self.decoder(z)
@@ -159,6 +156,26 @@ class CVAE(tf.keras.Model):
return x
def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask=None, noise=None):
if noise is None:
eps = tf.random.normal(shape=(uv_template.shape[0], self.latent_dim))
else:
eps = noise
z_cond = self.encode_conds(uv_template, uv_body_posed, uv_curv)
z = tf.concat([eps, z_cond], axis=-1)
x = self.decoder(z)
if self.avg_garment is not None:
x = x + self.avg_garment[tf.newaxis,:,:,:]
if uv_mask is not None:
x = x * uv_mask[:,:,:,tf.newaxis]
return x
def reparameterize(self, mean, logvar, fix_noise=False):
if not fix_noise:
eps = tf.random.normal(shape=mean.shape)
@@ -176,22 +193,37 @@ class CVAE(tf.keras.Model):
return logits
def call(self, inputs, training=None, zero_latent=False, fix_noise=False):
uv_unposed, uv_template, uv_body_posed, uv_curv = 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])
def encode_conds(self, uv_template, uv_body_posed, uv_curv):
if uv_template.shape[1] == 3:
uv_template = tf.transpose(uv_template, [0,2,3,1])
if uv_body_posed.shape[1] == 3:
uv_body_posed = tf.transpose(uv_body_posed, [0,2,3,1])
if uv_curv.shape[1] == 3:
uv_curv = tf.transpose(uv_curv, [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)
x = tf.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], axis=-1)
mean, logvar = self.encoder(x)
z = self.reparameterize(mean, logvar, fix_noise)
z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
z_cond = self.cond_encoder(z_cond)
if zero_latent:
z = tf.zeros_like(z)
return z_cond
def call(self, inputs, training=None, use_mask=False):
if use_mask:
uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask = inputs
else:
uv_unposed, uv_template, uv_body_posed, uv_curv = 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])
x = tf.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], axis=-1)
mean, logvar = self.encoder(x)
z = self.reparameterize(mean, logvar)
z_cond = self.encode_conds(uv_template, uv_body_posed, uv_curv)
z = tf.concat([z, z_cond], axis=-1)
@@ -200,14 +232,14 @@ class CVAE(tf.keras.Model):
# z = tf.concat([z, z_curv], axis=-1)
x = self.decoder(z)
# x = tf.transpose(x, [0,3,1,2])
# x = x * uv_mask[:,tf.newaxis,:,:]
# x = tf.image.resize(x, (512, 256), tf.image.ResizeMethod.LANCZOS5)
# x = x + tf.transpose(uv_shaped_body, [0,2,3,1])
if self.avg_garment is not None:
x = x + self.avg_garment[tf.newaxis,:,:,:]
# x = x * uv_mask[:,:,:,tf.newaxis]
if use_mask:
x = x * uv_mask[:,:,:,tf.newaxis]
return x, mean, logvar, z