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/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):