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

backup

Browse Source
This commit is contained in:
luni
2022-11-04 16:17:58 +00:00
parent ebe8522905
commit 1c6b805519
18 changed files with 1413 additions and 486 deletions
Download Patch File Download Diff File Expand all files Collapse all files
models/cvae_model.py
+23 -15
View File
@@ -26,14 +26,16 @@ class CVAE(tf.keras.Model):
if config['body_model'] != '':
load_model(self.body_enc, config['body_model'] + '_enc')
self.curv_enc : AEEncoder = AEEncoder(config['cond_ae']['latent_dim'], config['cond_ae']['encoder_depths'], config['cond_ae']['encoder_dims'], config['cond_ae']['bottleneck_size'], config['cond_ae']['kernel_size'])
if config['curvature_model'] != '':
load_model(self.curv_enc, config['curvature_model'] + '_enc')
if config['use_curvature']:
self.curv_enc : AEEncoder = AEEncoder(config['cond_ae']['latent_dim'], config['cond_ae']['encoder_depths'], config['cond_ae']['encoder_dims'], config['cond_ae']['bottleneck_size'], config['cond_ae']['kernel_size'])
if config['curvature_model'] != '':
load_model(self.curv_enc, config['curvature_model'] + '_enc')
if config['freeze_encoders']:
self.static_enc.trainable = False
self.body_enc.trainable = False
self.curv_enc.trainable = False
if config['use_curvature']:
self.curv_enc.trainable = False
self.cond_encoder = tf.keras.Sequential([tf.keras.layers.Dense(512, activation='relu'),
tf.keras.layers.Dense(256, activation='relu'),
@@ -50,7 +52,13 @@ class CVAE(tf.keras.Model):
self.add_noise = config['cvae']['add_noise']
self.noise_amount = config['cvae']['noise_amount']
self.use_curvature = config['use_curvature']
self.use_offsets_over_body = not config['use_skinning']
if self.use_offsets_over_body:
assert self.avg_garment is None, 'When we are not using skinning, we can\'t use avg_garment!'
if self.avg_garment is not None:
assert not self.use_offsets_over_body, 'When avg_garment is given, we must use skinning!'
@tf.function
def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask, noise=None):
@@ -84,27 +92,24 @@ class CVAE(tf.keras.Model):
def encode_conds(self, uv_template, uv_body_posed, uv_curv, training=None):
z_body = self.body_enc(uv_body_posed, training=training)
z_static = self.static_enc(uv_template, training=training)
z_curv = self.curv_enc(uv_curv, training=training)
if self.use_curvature:
z_curv = self.curv_enc(uv_curv, training=training)
z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
else:
z_cond = tf.concat([z_static, z_body], axis=-1)
z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
z_cond = self.cond_encoder(z_cond, training=training)
return z_cond
def call(self, inputs, training=None):
if self.size == 512:
uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask = inputs
uv_template_high, uv_body_posed_high, uv_curv_high = uv_template, uv_body_posed, uv_curv
else:
uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask, uv_template_high, uv_body_posed_high, uv_curv_high, uv_mask_high = inputs
uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask = inputs
x = tf.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], axis=-1)
mean, logvar = self.encoder(x, training=training)
z = self.reparameterize(mean, logvar)
z_cond = self.encode_conds(uv_template_high, uv_body_posed_high, uv_curv_high, training=training)
z_cond = self.encode_conds(uv_template, uv_body_posed, uv_curv, training=training)
if self.add_noise and training == True:
z_cond = z_cond + (tf.random.uniform(z_cond.shape, maxval=self.noise_amount * 2) - self.noise_amount)
@@ -115,7 +120,10 @@ class CVAE(tf.keras.Model):
if self.avg_garment is not None:
x = x + self.avg_garment[tf.newaxis,:,:,:]
if self.use_offsets_over_body:
x = x + uv_body_posed
x = x * uv_mask[:,:,:,tf.newaxis]
return x, mean, logvar
models/pyramid_model_optimized.py
+75 -31
View File
@@ -42,6 +42,28 @@ class Pyramid(tf.keras.Model):
self.levels = [PyramidLevel({ **config['pyramid']['levels'][i], **{ 'cond_ae': config['cond_ae']} }) for i in range(len(config['pyramid']['levels']))]
# self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
# self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
# self.levels[-3].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
if len(config['pyramid']['resolutions']) == 2:
self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 10)])
if len(config['pyramid']['resolutions']) == 3:
self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 10)])
if len(config['pyramid']['resolutions']) == 4:
self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 1000)])
self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
self.levels[-3].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 10)])
if len(config['pyramid']['resolutions']) == 5:
self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 10000)])
self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 1000)])
self.levels[-3].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
self.levels[-4].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 10)])
self.add_noise = config['pyramid']['add_noise']
self.use_offsets = config['use_offsets']
self.upscale_enhancer = None
if config['pyramid']['final_upscale_enhancer'] != '':
@@ -66,42 +88,64 @@ class Pyramid(tf.keras.Model):
def call(self, inputs, training=None):
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)
# out_levels = []
# pred, mean, logvar = self.levels[0](levels_input[0], z_cond, training=training)
# if 1 != len(self.levels):
# pred = self.levels[0].enhance(pred, levels_input[0][-2], training=training) # levels_input[i][-2]=uv_mask_res
# pred = tf.image.resize(pred, (self.levels[1].size, self.levels[1].size//2), tf.image.ResizeMethod.LANCZOS5)
# out_levels.append([pred, mean, logvar])
# pred = pred * levels_input[1][-2][:,:,:,tf.newaxis]
# else:
# pred = self.enhance(pred, levels_input[0][-2], training=training)
# pred = tf.image.resize(pred, (512, 256), tf.image.ResizeMethod.LANCZOS5)
# pred = pred * levels_input[-1][-1][:,:,:,tf.newaxis] # full res uv mask
# out_levels.append([pred, mean, logvar])
if self.add_noise:
z_cond = z_cond + tf.random.uniform(z_cond.shape, maxval=10)
# for i in range(1, len(self.levels)):
# pred_level, mean, logvar = self.levels[i](levels_input[i], z_cond, training=training)
# pred = pred + pred_level
# if i != len(self.levels) - 1:
# pred = self.levels[i].enhance(pred, levels_input[i][-2], training=training)
# out_levels.append([pred, mean, logvar])
# pred = tf.image.resize(pred, (self.levels[i+1].size, self.levels[i+1].size//2), tf.image.ResizeMethod.LANCZOS5)
# pred = pred * levels_input[i+1][-2][:,:,:,tf.newaxis]
# else:
# if self.upscale_enhancer is not None:
# pred = self.enhance(pred, levels_input[-1][-2], training=training)
# pred = tf.image.resize(pred, (512, 256), tf.image.ResizeMethod.LANCZOS5)
# pred = pred * levels_input[-1][-1][:,:,:,tf.newaxis] # full res uv mask
# out_levels.append([pred, mean, logvar])
out_levels = []
pred, mean, logvar = self.levels[0](levels_input[0], z_cond, training=training)
if self.use_offsets:
pred = pred + levels_input[0][2] # uv_body_posed_32
if 1 != len(self.levels):
pred = self.levels[0].enhance(pred, levels_input[0][-2], training=training) # levels_input[i][-2]=uv_mask_res
pred = tf.image.resize(pred, (self.levels[1].size, self.levels[1].size//2), tf.image.ResizeMethod.LANCZOS5)
out_levels.append([pred, mean, logvar])
pred = pred * levels_input[1][-2][:,:,:,tf.newaxis]
else:
pred = self.enhance(pred, levels_input[0][-2], training=training)
pred = tf.image.resize(pred, (512, 256), tf.image.ResizeMethod.LANCZOS5)
pred = pred * levels_input[-1][-1][:,:,:,tf.newaxis] # full res uv mask
out_levels.append([pred, mean, logvar])
pred_0, mean_0, logvar_0 = self.levels[0](levels_input[0], z_cond, training=training)
pred_0 = self.enhance(pred_0, levels_input[0][-2], training=training)
pred_0 = tf.image.resize(pred_0, (512, 256), tf.image.ResizeMethod.LANCZOS5)
pred_0 = pred_0 * levels_input[-1][-1][:,:,:,tf.newaxis] # full res uv mask
for i in range(1, len(self.levels)):
pred_level, mean, logvar = self.levels[i](levels_input[i], z_cond, training=training)
pred = pred + pred_level
if i != len(self.levels) - 1:
pred = self.levels[i].enhance(pred, levels_input[i][-2], training=training)
out_levels.append([pred, mean, logvar])
pred = tf.image.resize(pred, (self.levels[i+1].size, self.levels[i+1].size//2), tf.image.ResizeMethod.LANCZOS5)
pred = pred * levels_input[i+1][-2][:,:,:,tf.newaxis]
else:
if self.upscale_enhancer is not None:
pred = self.enhance(pred, levels_input[-1][-2], training=training)
pred = tf.image.resize(pred, (512, 256), tf.image.ResizeMethod.LANCZOS5)
pred = pred * levels_input[-1][-1][:,:,:,tf.newaxis] # full res uv mask
out_levels.append([pred, mean, logvar])
return [ [pred_0, mean_0, logvar_0] ]
return out_levels
# # 32
# pred_0, mean_0, logvar_0 = self.levels[0](levels_input[0], z_cond, training=training)
# pred_0 = self.enhance(pred_0, levels_input[0][-2], training=training)
# pred_0 = tf.image.resize(pred_0, (512, 256), tf.image.ResizeMethod.LANCZOS5)
# pred_0 = pred_0 * levels_input[-1][-1][:,:,:,tf.newaxis] # full res uv mask
# return [ [pred_0, mean_0, logvar_0] ]
# # 32 64
# pred_0, mean_0, logvar_0 = self.levels[0](levels_input[0], z_cond, training=training)
# pred_0 = self.levels[0].enhance(pred_0, levels_input[0][-2], training=training) # levels_input[i][-2]=uv_mask_res
# pred_0 = tf.image.resize(pred_0, (self.levels[1].size, self.levels[1].size//2), tf.image.ResizeMethod.LANCZOS5)
# pred = pred_0 * levels_input[1][-2][:,:,:,tf.newaxis]
# pred_1, mean_1, logvar_1 = self.levels[1](levels_input[1], z_cond, training=training)
# pred_1 = pred + pred_1
# pred_1 = self.enhance(pred_1, levels_input[-1][-2], training=training)
# pred_1 = tf.image.resize(pred_1, (512, 256), tf.image.ResizeMethod.LANCZOS5)
# pred_1 = pred_1 * levels_input[-1][-1][:,:,:,tf.newaxis] # full res uv mask
# return [ [pred_0, mean_0, logvar_0], [pred_1, mean_1, logvar_1] ]
# class Pyramid(tf.keras.Model):