resnet pyramid
This commit is contained in:
+2
-2
@@ -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,
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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
@@ -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
@@ -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):
|
||||
|
||||
Reference in New Issue
Block a user