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

update pyramid

Browse Source
This commit is contained in:
luni
2022-10-16 14:57:36 +00:00
parent a6834caf80
commit 07c9fb57df
3 changed files with 10 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
models/pyramid_level.py
+2 -2
View File
@@ -9,7 +9,7 @@ import scipy as sp
from custom_utils.custom_utils import load_model
from models.base import AEEncoder, VariationalEncoder, Decoder
from models.upscaler import UpscalerBasic, UpscalerBasicNew
from models.upscaler import UpscalerBasic, UpscalerBasicConv, UpscalerBasicNew
class PyramidLevel(tf.keras.Model):
@@ -24,7 +24,7 @@ class PyramidLevel(tf.keras.Model):
self.upscale_enhancer = None
if config['upscale_enhancer'] != '':
self.upscale_enhancer = UpscalerBasicNew(config['resolution'])
self.upscale_enhancer = UpscalerBasicConv(config['resolution'])
load_model(self.upscale_enhancer, config['upscale_enhancer'])
self.upscale_enhancer.trainable = not config['upscale_enhancer_freeze']
models/pyramid_model.py
+3 -3
View File
@@ -10,7 +10,7 @@ from custom_utils.custom_utils import load_model
from models.base import AEEncoder
from models.cvae_model import CVAE
from models.pyramid_level import PyramidLevel
from models.upscaler import UpscalerBasic, UpscalerBasicNew
from models.upscaler import UpscalerBasic, UpscalerBasicConv, UpscalerBasicNew
class Pyramid(tf.keras.Model):
@@ -22,7 +22,7 @@ class Pyramid(tf.keras.Model):
if config['pyramid']['static_model'] != '':
load_model(self.static_enc, config['pyramid']['static_model'] + '_enc')
self.body_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'])
self.body_enc : AEEncoder = AEEncoder(config['cond_ae_body']['latent_dim'], config['cond_ae_body']['encoder_depths'], config['cond_ae_body']['encoder_dims'], config['cond_ae_body']['bottleneck_size'], config['cond_ae_body']['kernel_size'])
if config['pyramid']['body_model'] != '':
load_model(self.body_enc, config['pyramid']['body_model'] + '_enc')
@@ -45,7 +45,7 @@ class Pyramid(tf.keras.Model):
self.upscale_enhancer = None
if config['pyramid']['final_upscale_enhancer'] != '':
self.upscale_enhancer = UpscalerBasicNew(low_res=config['pyramid']['resolutions'][-1])
self.upscale_enhancer = UpscalerBasicConv(low_res=config['pyramid']['resolutions'][-1])
load_model(self.upscale_enhancer, config['pyramid']['final_upscale_enhancer'])
self.upscale_enhancer.trainable = False
models/upscaler.py
+5 -5
View File
@@ -125,19 +125,19 @@ class UpscalerBasicNew(tf.keras.Model):
class UpscalerBasicConv(tf.keras.Model):
def __init__(self, low_res=32):
def __init__(self, low_res=32, kernel_size=3):
super(UpscalerBasicConv, self).__init__()
self.low_res = low_res
self.seq = tf.keras.Sequential([
tf.keras.layers.Conv2D(4096, kernel_size=3, padding='same', activation='relu', name='upscale_enhancer_conv_1'),
tf.keras.layers.Conv2D(256, 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(2048, kernel_size=3, padding='same', activation='relu', name='upscale_enhancer_conv_2'),
tf.keras.layers.Conv2D(128, 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(1024, kernel_size=3, padding='same', activation='relu', name='upscale_enhancer_conv_3'),
tf.keras.layers.Conv2D(64, 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(128, kernel_size=3, padding='same', activation='relu', name='upscale_enhancer_conv_4'),
tf.keras.layers.Conv2D(32, 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'),
])