backup

2022-10-04 21:50:58 +00:00
parent 5b82041acb
commit 3c798b473d
8 changed files with 270 additions and 125 deletions
@@ -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

@@ -10,6 +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


 class Pyramid(tf.keras.Model):
@@ -41,6 +42,16 @@ class Pyramid(tf.keras.Model):

        self.levels = [PyramidLevel({ **config['pyramid']['levels'][i], **{ 'cond_ae': config['cond_ae']} }) for i in range(len(config['pyramid']['levels']))]

+ 
+        self.upscale_enhancer = None
+        if config['pyramid']['final_upscale_enhancer'] != '':
+            self.upscale_enhancer = UpscalerBasicNew(low_res=config['pyramid']['resolutions'][-1])
+            load_model(self.upscale_enhancer, config['pyramid']['final_upscale_enhancer'])
+            self.upscale_enhancer.trainable = False
+
+    def enhance(self, low, mask, training=None):
+        x = self.upscale_enhancer([low, mask], training=training)
+        return x

    def encode_conds(self, uv_template, uv_body_posed, uv_curv):
        z_static = self.static_enc(uv_template)
@@ -53,18 +64,37 @@ class Pyramid(tf.keras.Model):
        return z_cond

    def call(self, inputs, training=None):
-        uv_static, uv_body, uv_curv = inputs[-1][1], inputs[-1][2], inputs[-1][3]
+        levels_input, conds = inputs
+        uv_static, uv_body, uv_curv = conds[0], conds[1], conds[2]
        z_cond = self.encode_conds(uv_static, uv_body, uv_curv)

        out_levels = []
-        pred, mean, logvar = self.levels[0](inputs[0], z_cond, training=training)
-        out_levels.append([pred, mean, logvar])
+        pred, mean, logvar = self.levels[0](levels_input[0], z_cond, training=training)
+        pred = self.levels[0].enhance(pred, levels_input[0][-2], training=training)  # levels_input[i][-2]=uv_mask_res
+        if 1 != len(self.levels):
+            pred = tf.image.resize(pred, (self.levels[1].size, self.levels[1].size//2), tf.image.ResizeMethod.LANCZOS5)
+            out_levels.append([pred, mean, logvar])
+            pred = pred * levels_input[1][-2][:,:,:,tf.newaxis]
+        else:
+            pred = tf.image.resize(pred, (512, 256), tf.image.ResizeMethod.LANCZOS5)
+            out_levels.append([pred, mean, logvar])

        for i in range(1, len(self.levels)):
-            pred_level, mean, logvar = self.levels[i](inputs[i], z_cond, training=training)
+            pred_level, mean, logvar = self.levels[i](levels_input[i], z_cond, training=training)
            pred = pred + pred_level
-            out_levels.append([pred, mean, logvar])
-        
+            if i != len(self.levels) - 1:
+                pred = self.levels[i].enhance(pred, levels_input[i][-2], training=training)
+                out_levels.append([pred, mean, logvar])
+                pred = tf.image.resize(pred, (self.levels[i+1].size, self.levels[i+1].size//2), tf.image.ResizeMethod.LANCZOS5)
+                pred = pred * levels_input[i+1][-2][:,:,:,tf.newaxis]
+            else:
+                if self.upscale_enhancer is not None:
+                    pred = self.enhance(pred, levels_input[-1][-2], training=training)
+                    pred = tf.image.resize(pred, (512, 256), tf.image.ResizeMethod.LANCZOS5)
+                out_levels.append([pred, mean, logvar])
+                pred = pred * levels_input[-1][-1][:,:,:,tf.newaxis] # full res uv mask
+
+
        return out_levels


@@ -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):