backup

models/upscaler.py

@@ -55,9 +55,11 @@ class Upscaler(tf.keras.Model):
 
 class UpscalerBasic(tf.keras.Model):
     
-    def __init__(self, bottleneck_size, size=32):
+    def __init__(self, bottleneck_size, low_res=32, high_res=64):
         super(UpscalerBasic, self).__init__()
 
+        self.low_res = low_res
+        self.high_res = high_res
         self.flatten = tf.keras.layers.Flatten()
 
         self.seq = tf.keras.Sequential([
@@ -67,10 +69,10 @@ class UpscalerBasic(tf.keras.Model):
                     tf.keras.layers.BatchNormalization(),
                     tf.keras.layers.Dense(1024, activation='relu'),
                     tf.keras.layers.BatchNormalization(),
-                    tf.keras.layers.Dense(size * (size//2) * 3),
+                    tf.keras.layers.Dense(low_res * (low_res//2) * 3),
         ])
 
-        self.reshape = tf.keras.layers.Reshape((size, (size//2), 3))
+        self.reshape = tf.keras.layers.Reshape((low_res, (low_res//2), 3))
 
     def call(self, inputs, training=None):
         img, mask = inputs
@@ -81,10 +83,46 @@ class UpscalerBasic(tf.keras.Model):
 
         x = x * (1 - mask[:,:,:,tf.newaxis])
         x = x + img
-        x = tf.image.resize(x, size=(512, 256), method=tf.image.ResizeMethod.LANCZOS5)
+        x = tf.image.resize(x, size=(self.high_res, self.high_res//2), method=tf.image.ResizeMethod.LANCZOS5)
 
         return x
 
+
+
+
+class UpscalerBasicNew(tf.keras.Model):
+    
+    def __init__(self, low_res=32):
+        super(UpscalerBasicNew, self).__init__()
+
+        self.low_res = low_res
+        self.flatten = tf.keras.layers.Flatten(name='upscale_enhancer_flattent')
+
+        self.seq = tf.keras.Sequential([
+                    tf.keras.layers.Dense(1024, activation='relu', name='upscale_enhancer_dense_1'),
+                    tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_1'),
+                    tf.keras.layers.Dense(1024, activation='relu', name='upscale_enhancer_dense_2'),
+                    tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_2'),
+                    tf.keras.layers.Dense(1024, activation='relu', name='upscale_enhancer_dense_3'),
+                    tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_3'),
+                    tf.keras.layers.Dense(low_res * (low_res//2) * 3, activation='tanh', name='upscale_enhancer_dense_final'),
+        ])
+
+        self.reshape = tf.keras.layers.Reshape((low_res, (low_res//2), 3), name='upscale_enhancer_reshape')
+
+    def call(self, inputs, training=None):
+        img, mask = inputs
+
+        x = self.flatten(img)
+        x = self.seq(x, training=training)
+        x = self.reshape(x)
+
+        x = x * (1 - mask[:,:,:,tf.newaxis])
+        x = x + img
+
+        return x
+
+
 # class Upscaler(tf.keras.Model):
     
 #     def __init__(self, bottleneck_size):