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

save

Browse Source
This commit is contained in:
luni
2023-03-07 22:02:55 +00:00
parent cbecf668df
commit cb76d36c4f
20 changed files with 626 additions and 168 deletions
Download Patch File Download Diff File Expand all files Collapse all files
models/pyramid_level.py
+3 -3
View File
@@ -101,6 +101,7 @@ class PyramidLevelNew(tf.keras.Model):
def __init__(self, config):
super(PyramidLevelNew, self).__init__()
self.size = config['resolution']
self.latent_dim=config['latent_dim']
self.encoder = VariationalEncoderNew(config['encoder_depths'], config['encoder_dims'], kernel_size=config['kernel_size'], latent_dim=config['latent_dim'])
self.decoder = DecoderNew(config['decoder_depths'], config['decoder_dims'], config['kernel_size'])
@@ -126,13 +127,12 @@ class PyramidLevelNew(tf.keras.Model):
def sample(self, z_cond, uv_mask, noise=None):
if noise is None:
eps = tf.random.normal(shape=(uv_mask.shape[0], self.latent_dim))
eps = tf.random.normal(shape=(uv_mask.shape[0], 8, 4, self.latent_dim))
else:
eps = noise
z = tf.concat([eps, z_cond], axis=-1)
x = self.decoder(z)
x = self.decoder(z, training=False)
if self.avg_garment is not None:
x = x + self.avg_garment[tf.newaxis,:,:,:]
models/pyramid_model_optimized.py
+81
View File
@@ -320,6 +320,87 @@ class PyramidNew(tf.keras.Model):
return z_cond
def sample(self, inputs):
levels_input, conds = inputs
uv_static, uv_body, uv_curv = conds[0], conds[1], conds[2]
z_cond = self.encode_conds(uv_static, uv_body, uv_curv, training=False)
pred = [None, None, None, None, None]
for i in range(len(self.levels)):
pred[i] = self.levels[i].sample(z_cond, levels_input[i][-2])
pred_1 = pred[0]
pred_2 = pred[1]
pred_3 = pred[2]
pred_4 = pred[3]
pred_5 = pred[4]
pred_1 = self.levels[0].enhance(pred_1, levels_input[0][-2], training=False)
pred_1 = tf.image.resize(pred_1, (self.levels[1].size, self.levels[1].size//2), tf.image.ResizeMethod.BILINEAR)
pred_1 = pred_1 * levels_input[1][-2][:,:,:,tf.newaxis]
pred_2 = pred_1 + pred_2
pred_2 = self.levels[1].enhance(pred_2, levels_input[1][-2], training=False)
pred_2 = tf.image.resize(pred_2, (self.levels[2].size, self.levels[2].size//2), tf.image.ResizeMethod.BILINEAR)
pred_2 = pred_2 * levels_input[2][-2][:,:,:,tf.newaxis]
pred_3 = pred_2 + pred_3
pred_3 = self.levels[2].enhance(pred_3, levels_input[2][-2], training=False)
pred_3 = tf.image.resize(pred_3, (self.levels[3].size, self.levels[3].size//2), tf.image.ResizeMethod.BILINEAR)
pred_3 = pred_3 * levels_input[3][-2][:,:,:,tf.newaxis]
pred_4 = pred_3 + pred_4
pred_4 = self.levels[3].enhance(pred_4, levels_input[3][-2], training=False)
pred_4 = tf.image.resize(pred_4, (self.levels[4].size, self.levels[4].size//2), tf.image.ResizeMethod.BILINEAR)
pred_4 = pred_4 * levels_input[4][-2][:,:,:,tf.newaxis]
pred_5 = pred_4 + pred_5
pred_5 = pred_5 * levels_input[4][-2][:,:,:,tf.newaxis]
return pred_5
def sample_seq(self, inputs):
levels_input, conds = inputs
uv_static, uv_body, uv_curv = conds[0], conds[1], conds[2]
z_cond = self.encode_conds(uv_static, uv_body, uv_curv, training=False)
# level 1
pred_1 = self.levels[0].sample(z_cond, levels_input[0][-2])
pred_1 = self.levels[0].enhance(pred_1, levels_input[0][-2], training=False)
pred_1 = tf.image.resize(pred_1, (self.levels[1].size, self.levels[1].size//2), tf.image.ResizeMethod.BILINEAR)
pred_1 = pred_1 * levels_input[1][-2][:,:,:,tf.newaxis]
# level 2
pred_2 = self.levels[1].sample(z_cond, levels_input[1][-2])
pred_2 = pred_1 + pred_2
pred_2 = self.levels[1].enhance(pred_2, levels_input[1][-2], training=False)
pred_2 = tf.image.resize(pred_2, (self.levels[2].size, self.levels[2].size//2), tf.image.ResizeMethod.BILINEAR)
pred_2 = pred_2 * levels_input[2][-2][:,:,:,tf.newaxis]
# level 3
pred_3 = self.levels[2].sample(z_cond, levels_input[2][-2])
pred_3 = pred_2 + pred_3
pred_3 = self.levels[2].enhance(pred_3, levels_input[2][-2], training=False)
pred_3 = tf.image.resize(pred_3, (self.levels[3].size, self.levels[3].size//2), tf.image.ResizeMethod.BILINEAR)
pred_3 = pred_3 * levels_input[3][-2][:,:,:,tf.newaxis]
# level 4
pred_4 = self.levels[3].sample(z_cond, levels_input[3][-2])
pred_4 = pred_3 + pred_4
pred_4 = self.levels[3].enhance(pred_4, levels_input[3][-2], training=False)
pred_4 = tf.image.resize(pred_4, (self.levels[4].size, self.levels[4].size//2), tf.image.ResizeMethod.BILINEAR)
pred_4 = pred_4 * levels_input[4][-2][:,:,:,tf.newaxis]
# level 5
pred_5 = self.levels[4].sample(z_cond, levels_input[4][-2])
pred_5 = pred_5 * levels_input[4][-2][:,:,:,tf.newaxis]
pred_5 = pred_4 + pred_5
return pred_5
def call(self, inputs, training):
levels_input, conds = inputs
uv_static, uv_body, uv_curv = conds[0], conds[1], conds[2]