backup

models/pyramid_level.py

@@ -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
+from models.upscaler import UpscalerBasic, UpscalerBasicNew
 
 
 class PyramidLevel(tf.keras.Model):
@@ -24,13 +24,9 @@ class PyramidLevel(tf.keras.Model):
 
         self.upscale_enhancer = None
         if config['upscale_enhancer'] != '':
-            self.upscale_enhancer = UpscalerBasic(512, config['resolution'])
+            self.upscale_enhancer = UpscalerBasicNew(config['resolution'])
             load_model(self.upscale_enhancer, config['upscale_enhancer'])
-            self.upscale_enhancer = tf.keras.Sequential([
-                self.upscale_enhancer.flatten,
-                self.upscale_enhancer.seq,
-                self.upscale_enhancer.reshape
-                ])
+            self.upscale_enhancer.trainable = not config['upscale_enhancer_freeze']
 
         if config['avg_garment_path'] != '':
             self.avg_garment = np.load(config['avg_garment_path'])
@@ -39,10 +35,8 @@ class PyramidLevel(tf.keras.Model):
 
 
     def enhance(self, low, mask, training=None):
-        x = self.upscale_enhancer(low, training=training)
-        x = x * (1 - mask[:,:,:,tf.newaxis])
-        low = x + low
-        return low
+        x = self.upscale_enhancer([low, mask], training=training)
+        return x
 
 
     def sample(self, z_cond, uv_mask, noise=None):
@@ -87,11 +81,11 @@ class PyramidLevel(tf.keras.Model):
 
         x = x * uv_mask[:,:,:,tf.newaxis]
 
-        if self.upscale_enhancer is not None:
-            pred_low = self.enhance(x, uv_mask, training=training)
-            x = tf.image.resize(pred_low, (512, 256), tf.image.ResizeMethod.LANCZOS5)
+        # if self.upscale_enhancer is not None:
+        #     pred_low = self.enhance(x, uv_mask, training=training)
+        #     x = tf.image.resize(pred_low, (512, 256), tf.image.ResizeMethod.LANCZOS5)
         
-        x = x * uv_mask_high[:,:,:,tf.newaxis]
+        # x = x * uv_mask_high[:,:,:,tf.newaxis]
 
         return x, mean, logvar