resnet pyramid

This commit is contained in:
2023-03-02 17:40:14 +00:00
parent 532f5e12f4
commit ec51e5e7cb
7 changed files with 261 additions and 64 deletions
+2 -2
View File
@@ -113,10 +113,10 @@ class Decoder(tf.keras.Model):
class VariationalEncoderNew(tf.keras.Model):
def __init__(self, features):
def __init__(self, layer_params, features):
super(VariationalEncoderNew, self).__init__()
self.encoder = resnet_18(features)
self.encoder = resnet_18(layer_params, features)
self.mean_branch = tf.keras.Sequential([
tf.keras.layers.Conv2D(filters=features,
+3 -3
View File
@@ -102,8 +102,8 @@ class PyramidLevelNew(tf.keras.Model):
super(PyramidLevelNew, self).__init__()
self.size = config['resolution']
self.encoder = VariationalEncoderNew(config['encoder_depths'], config['encoder_dims'], kernel_size=config['kernel_size'])
self.decoder = DecoderNew(config['decoder_depths'], config['decoder_dims'], config['kernel_size'])
self.encoder = VariationalEncoderNew(config['layer_params'], config['bottleneck_size'])
self.decoder = DecoderNew(config['layer_params'], config['bottleneck_size'])
self.upscale_enhancer = None
if config['upscale_enhancer'] != '':
@@ -117,7 +117,7 @@ class PyramidLevelNew(tf.keras.Model):
else:
self.avg_garment = None
self.reshape = tf.keras.layers.Reshape((8, 4, 64))
self.reshape = tf.keras.layers.Reshape((8, 4, config['bottleneck_size']))
def enhance(self, low, mask, training=None):
+3 -3
View File
@@ -256,15 +256,15 @@ class PyramidNew(tf.keras.Model):
def __init__(self, config):
super(PyramidNew, self).__init__()
self.static_enc : AEEncoderNew = AEEncoderNew(config['cond_ae_static']['encoder_depths'], config['cond_ae_static']['encoder_dims'], config['cond_ae_static']['kernel_size'], config['cond_ae_static']['activation'])
self.static_enc : AEEncoderNew = AEEncoderNew(config['cond_ae_static']['layer_params'], config['cond_ae_static']['bottleneck_size'], config['cond_ae_static']['activation'])
if config['pyramid']['static_model'] != '':
load_model(self.static_enc, config['pyramid']['static_model'] + '_enc')
self.body_enc : AEEncoderNew = AEEncoderNew(config['cond_ae_body']['encoder_depths'], config['cond_ae_body']['encoder_dims'], config['cond_ae_body']['kernel_size'])
self.body_enc : AEEncoderNew = AEEncoderNew(config['cond_ae_body']['layer_params'], config['cond_ae_body']['bottleneck_size'], config['cond_ae_static']['activation'])
if config['pyramid']['body_model'] != '':
load_model(self.body_enc, config['pyramid']['body_model'] + '_enc')
self.curv_enc : AEEncoderNew = AEEncoderNew(config['cond_ae_curv']['encoder_depths'], config['cond_ae_curv']['encoder_dims'], config['cond_ae_curv']['kernel_size'], config['cond_ae_curv']['activation'])
self.curv_enc : AEEncoderNew = AEEncoderNew(config['cond_ae_curv']['layer_params'], config['cond_ae_curv']['bottleneck_size'], config['cond_ae_curv']['activation'])
if config['pyramid']['curvature_model'] != '':
load_model(self.curv_enc, config['pyramid']['curvature_model'] + '_enc')
+42 -27
View File
@@ -131,9 +131,10 @@ class ResNetTypeI(tf.keras.Model):
super(ResNetTypeI, self).__init__()
self.reversed = reversed
self.labels = labels
self.layer_params = layer_params
if not reversed:
self.conv1 = tf.keras.layers.Conv2D(filters=32,
self.conv1 = tf.keras.layers.Conv2D(filters=32 if len(layer_params) > 1 else features,
kernel_size=(7, 7),
strides=2,
padding="same")
@@ -144,31 +145,36 @@ class ResNetTypeI(tf.keras.Model):
else:
self.conv1 = tf.keras.layers.Conv2DTranspose(filters=3 if labels is None else labels,
kernel_size=(7, 7),
strides=2,
strides=2 if len(layer_params) > 1 else 1,
padding="same")
self.bn1 = tf.keras.layers.BatchNormalization()
self.pool1 = tf.keras.layers.UpSampling2D()
self.layer1 = make_basic_block_layer(filter_num=32,
self.layer1 = make_basic_block_layer(filter_num=32 if len(layer_params) > 1 else (features if not is_cvae else features*2),
blocks=layer_params[0],
stride=1 if not reversed else 2,
upsample=reversed)
self.layer2 = make_basic_block_layer(filter_num=32,
blocks=layer_params[1],
stride=2,
upsample=reversed)
self.layer3 = make_basic_block_layer(filter_num=32,
blocks=layer_params[2],
stride=2,
upsample=reversed)
self.layer4 = make_basic_block_layer(filter_num=64,
blocks=layer_params[3],
stride=2,
upsample=reversed)
self.layer5 = make_basic_block_layer(filter_num=features if not is_cvae else features*2,
blocks=layer_params[4],
stride=2 if not reversed else 1,
upsample=reversed)
if len(layer_params) >= 2:
self.layer2 = make_basic_block_layer(filter_num=32,
blocks=layer_params[1],
stride=2,
upsample=reversed)
if len(layer_params) >= 3:
self.layer3 = make_basic_block_layer(filter_num=32,
blocks=layer_params[2],
stride=2,
upsample=reversed)
if len(layer_params) >= 4:
self.layer4 = make_basic_block_layer(filter_num=64,
blocks=layer_params[3],
stride=2,
upsample=reversed)
if len(layer_params) >= 5:
self.layer5 = make_basic_block_layer(filter_num=features if not is_cvae else features*2,
blocks=layer_params[4],
stride=2 if not reversed else 1,
upsample=reversed)
def call(self, inputs, training=None, mask=None):
if not self.reversed:
@@ -177,15 +183,24 @@ class ResNetTypeI(tf.keras.Model):
x = tf.nn.relu(x)
x = self.pool1(x)
x = self.layer1(x, training=training)
x = self.layer2(x, training=training)
x = self.layer3(x, training=training)
x = self.layer4(x, training=training)
x = self.layer5(x, training=training)
if len(self.layer_params) >= 2:
x = self.layer2(x, training=training)
if len(self.layer_params) >= 3:
x = self.layer3(x, training=training)
if len(self.layer_params) >= 4:
x = self.layer4(x, training=training)
if len(self.layer_params) >= 5:
x = self.layer5(x, training=training)
else:
x = self.layer5(inputs, training=training)
x = self.layer4(x, training=training)
x = self.layer3(x, training=training)
x = self.layer2(x, training=training)
x = inputs
if len(self.layer_params) >= 5:
x = self.layer5(x, training=training)
if len(self.layer_params) >= 4:
x = self.layer4(x, training=training)
if len(self.layer_params) >= 2:
x = self.layer3(x, training=training)
if len(self.layer_params) >= 3:
x = self.layer2(x, training=training)
x = self.layer1(x, training=training)
x = self.pool1(x)
+5 -5
View File
@@ -131,16 +131,16 @@ class UpscalerBasicConv(tf.keras.Model):
self.low_res = low_res
self.seq = tf.keras.Sequential([
tf.keras.layers.Conv2D(256, kernel_size=kernel_size, padding='same', activation='relu', name='upscale_enhancer_conv_1'),
tf.keras.layers.Conv2D(64, kernel_size=kernel_size, padding='same', activation='relu', name='upscale_enhancer_conv_1'),
tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_1'),
tf.keras.layers.Conv2D(128, kernel_size=kernel_size, padding='same', activation='relu', name='upscale_enhancer_conv_2'),
tf.keras.layers.Conv2D(64, kernel_size=kernel_size, padding='same', activation='relu', name='upscale_enhancer_conv_2'),
tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_2'),
tf.keras.layers.Conv2D(64, kernel_size=kernel_size, padding='same', activation='relu', name='upscale_enhancer_conv_3'),
tf.keras.layers.Conv2D(16, kernel_size=kernel_size, padding='same', activation='relu', name='upscale_enhancer_conv_3'),
tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_3'),
tf.keras.layers.Conv2D(32, kernel_size=kernel_size, padding='same', activation='relu', name='upscale_enhancer_conv_4'),
tf.keras.layers.Conv2D(16, kernel_size=kernel_size, padding='same', activation='relu', name='upscale_enhancer_conv_4'),
tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_4'),
tf.keras.layers.Dense(3, activation='tanh', name='upscale_enhancer_dense_final'),
])
])
def call(self, inputs, training):