commit all

models/ae_body_model.py

@@ -44,7 +44,7 @@ class AutoEncoder(tf.keras.Model):
         self.decoder = Decoder(decoder_input, decoder_depths, decoder_dims, kernel_size=kernel_size)
 
 
-    def call(self, inputs, training=None, use_mask=True):
+    def call(self, inputs, training, use_mask=True):
         if use_mask:
             x, mask = inputs
         else:

models/base.py

@@ -14,12 +14,12 @@ class Bottleneck(tf.keras.Model):
         self.flatten = tf.keras.layers.Flatten()
         self.dense = tf.keras.layers.Dense(latent_dim)
 
-    def call(self, inputs, training=None):
+    def call(self, inputs, training):
         x = inputs
 
         x = self.bottleneck(x, training=training)
         x = self.flatten(x)
-        x = self.dense(x)
+        x = self.dense(x, training=training)
 
         return x
 
@@ -35,7 +35,7 @@ class VariationalEncoder(tf.keras.Model):
         self.mean = tf.keras.layers.Dense(latent_dim)
         self.log_var = tf.keras.layers.Dense(latent_dim)
 
-    def call(self, inputs, training=None):
+    def call(self, inputs, training):
         x = inputs
 
         x = self.encoder(x, training=training)
@@ -55,7 +55,7 @@ class AEEncoder(tf.keras.Model):
         self.bottleneck = Bottleneck(bottleneck_size, latent_dim, kernel_size)
 
 
-    def call(self, inputs, training=None):
+    def call(self, inputs, training):
         if len(inputs) == 2:
             x, _ = inputs
         else:
@@ -85,7 +85,7 @@ class Decoder(tf.keras.Model):
             self.decoder = ConvNeXtReverseLabels(depths=list(reversed(depths)), dims=list(reversed(dims)), nr_of_classes=7)
             
 
-    def call(self, inputs, training=None):
+    def call(self, inputs, training):
         if len(inputs) == 2:
             x, mask = inputs
         else:
@@ -99,6 +99,105 @@ class Decoder(tf.keras.Model):
             x = tf.image.resize(x, size=(512, 256), method=tf.image.ResizeMethod.LANCZOS5)
 
 
+        if len(inputs) == 2:
+            if len(mask.shape) == 2:
+                x = x * mask[tf.newaxis,:,:,tf.newaxis]
+            else:
+                x = x * mask[:,:,:,tf.newaxis]
+
+        return x
+
+
+
+class VariationalEncoderNew(tf.keras.Model):
+    
+    def __init__(self, depths, dims, kernel_size=7):
+        super(VariationalEncoderNew, self).__init__()
+
+        self.encoder = ConvNeXt(depths=depths, dims=dims, include_top=False, kernel_size=kernel_size)
+
+        self.mean_branch = tf.keras.Sequential(
+                [Block(dim=64, drop_path=0, kernel_size=kernel_size,
+                       layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
+            )
+
+        self.var_branch = tf.keras.Sequential(
+                [Block(dim=64, drop_path=0, kernel_size=kernel_size,
+                       layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
+            )
+
+        self.flatten_mean = tf.keras.layers.Flatten()
+        self.flatten_var = tf.keras.layers.Flatten()
+
+
+
+    def call(self, inputs, training):
+        x = inputs
+
+        x = self.encoder(x, training=training)
+        mean = self.mean_branch(x)
+        log_var = self.var_branch(x)
+
+        mean = self.flatten_mean(mean)
+        log_var = self.flatten_var(log_var)
+
+        return mean, log_var
+        
+
+
+class AEEncoderNew(tf.keras.Model):
+    
+    def __init__(self, depths, dims, kernel_size=7, activation='tanh'):
+        super(AEEncoderNew, self).__init__()
+
+        self.encoder = ConvNeXt(depths=depths, dims=dims, include_top=False, kernel_size=kernel_size)
+        self.act = tf.keras.layers.Activation(activation)
+
+
+    def call(self, inputs, training):
+        if len(inputs) == 2:
+            x, _ = inputs
+        else:
+            x = inputs
+
+        x = self.encoder(x, training=training)
+        x = self.act(x)
+
+        if len(inputs) == 2:
+            return x, inputs[1]
+        else:
+            return x
+
+
+class DecoderNew(tf.keras.Model):
+    
+    def __init__(self, depths, dims, kernel_size=7, is_label=False):
+        super(DecoderNew, self).__init__()
+
+        self.is_label = is_label
+
+        if not is_label:
+            self.decoder = ConvNeXtReverse(depths=depths, dims=dims, kernel_size=kernel_size)
+        else:
+            self.decoder = ConvNeXtReverseLabels(depths=depths, dims=dims, nr_of_classes=7)
+            
+
+    def call(self, inputs, training):
+        if len(inputs) == 2:
+            x, mask = inputs
+        else:
+            x = inputs
+
+        x = self.decoder(x, training=training)
+
+        if self.is_label:
+            x = x * mask[:,:,:,tf.newaxis]
+
+            background = tf.ones((mask.shape[0],512,256,1)) * (1 - mask[:,:,:,tf.newaxis])
+            x = tf.concat((background, x), axis=-1)
+
+            x = tf.math.softmax(x)
+
         if len(inputs) == 2:
             if len(mask.shape) == 2:
                 x = x * mask[tf.newaxis,:,:,tf.newaxis]

models/cond_encoder.py

@@ -1,5 +1,5 @@
 import tensorflow as tf
-from models.base import AEEncoder, Decoder
+from models.base import AEEncoder, Decoder, AEEncoderNew, DecoderNew
 
 
 class AutoEncoder(tf.keras.Model):
@@ -11,10 +11,28 @@ class AutoEncoder(tf.keras.Model):
         self.decoder : Decoder = Decoder(config['decoder_input'], config['decoder_depths'], config['decoder_dims'], kernel_size=config['kernel_size'], is_label=config['is_label'], upscale=config['upscale_output'])
 
 
-    def call(self, inputs, training=None):
+    def call(self, inputs, training):
         x = inputs
 
         x = self.encoder(x, training=training)
         x = self.decoder(x, training=training)
 
         return x
+
+
+class AutoEncoderNew(tf.keras.Model):
+    
+    def __init__(self, config):
+        super(AutoEncoderNew, self).__init__()
+
+        self.encoder : AEEncoderNew = AEEncoderNew(config['encoder_depths'], config['encoder_dims'], config['kernel_size'], config['activation'])
+        self.decoder : DecoderNew = DecoderNew(config['decoder_depths'], config['decoder_dims'], kernel_size=config['kernel_size'], is_label=config['is_label'])
+
+
+    def call(self, inputs, training):
+        x = inputs
+
+        x = self.encoder(x, training=training)
+        x = self.decoder(x, training=training)
+
+        return x

models/cvae_curv_model.py

@@ -8,96 +8,43 @@ import scipy as sp
 # import plotly.express as ex
 
 from custom_utils.custom_utils import load_model
-from convnext.models.convnext_tf import ConvNeXt, Block, ConvNeXtReverseLabels
-from auto_encoder_body import AutoEncoder as BodyAE
-from auto_encoder_static import AutoEncoder as StaticAE
+from models.base import AEEncoder, VariationalEncoder, Decoder
 
 
 
-class Bottleneck(tf.keras.Model):
-    
-    def __init__(self, dim):
-        super(Bottleneck, self).__init__()
 
-        self.bottleneck = tf.keras.Sequential(
-                [Block(dim=dim, drop_path=0,
-                       layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
-            )
-
-    def call(self, inputs, training=None):
-        x = inputs
-
-        x = self.bottleneck(x, training=training)
-        
-        return x
-
-
-class Encoder(tf.keras.Model):
-    
-    def __init__(self, latent_dim):
-        super(Encoder, self).__init__()
-
-        self.encoder = ConvNeXt(depths=[3, 3, 9, 3], dims=[32, 64, 128, 256], include_top=False)
-        self.flatten = tf.keras.layers.Flatten()
-        self.mean = tf.keras.layers.Dense(latent_dim)
-        self.log_var = tf.keras.layers.Dense(latent_dim)
-
-    def call(self, inputs, training=None):
-        x = inputs
-
-        x = self.encoder(x)
-        x = self.flatten(x)
-        mean = self.mean(x)
-        log_var = self.log_var(x)
-
-        return mean, log_var
-
-
-class Decoder(tf.keras.Model):
-    
-    def __init__(self):
-        super(Decoder, self).__init__()
-
-        self.dense = tf.keras.layers.Dense(16 * 8 * 80)
-        self.reshape = tf.keras.layers.Reshape((16, 8, 80))
-        self.decoder = ConvNeXtReverseLabels(depths=list(reversed([3, 3, 9, 3])), dims=list(reversed([256, 128, 64, 16])))
-
-
-    def call(self, inputs, training=None):
-        x = inputs
-
-        x = self.dense(x)
-        x = self.reshape(x)
-        x = self.decoder(x)
-
-        return x
 
 
 class CVAE(tf.keras.Model):
     
-    def __init__(self, latent_dim=128, static_model='checkpoints/ae_static_3',
-                 body_model='checpoints/ae_body_6',):
+    def __init__(self, config):
         super(CVAE, self).__init__()
-        self.latent_dim = latent_dim
+        self.latent_dim = config['latent_dim']
+        self.latent_dim = config['cvae']['latent_dim']
+        self.size = config['resolution']
 
-        model : StaticAE = StaticAE(size=512)
-        load_model(model, static_model)
-        self.static_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
-        self.static_enc.trainable = False
+        self.static_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'])
+        if config['static_model'] != '':
+            load_model(self.static_enc, config['static_model'] + '_enc')
 
-        model : BodyAE = BodyAE(size=512)
-        load_model(model, body_model)
-        self.body_enc = tf.keras.Sequential([model.encoder, model.bottleneck])
-        self.body_enc.trainable = False
+        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['body_model'] != '':
+            load_model(self.body_enc, config['body_model'] + '_enc')
 
-        self.cond_encoder = tf.keras.Sequential([tf.keras.layers.Dense(256, activation='relu'),
+        if config['freeze_encoders']:
+            self.static_enc.trainable = False
+            self.body_enc.trainable = False
+
+        self.cond_encoder = tf.keras.Sequential([tf.keras.layers.Dense(512, activation='relu'),
                                                 tf.keras.layers.Dense(256, activation='relu'),
                                                 tf.keras.layers.Dense(128, activation='relu'),
-                                                tf.keras.layers.Dense(latent_dim),])
+                                                tf.keras.layers.Dense(self.latent_dim, activation='tanh'),])
 
-        self.encoder = Encoder(latent_dim)
-        self.decoder = Decoder()
+        self.encoder = VariationalEncoder(config['cvae']['latent_dim'], config['cvae']['encoder_depths'], config['cvae']['encoder_dims'], kernel_size=config['cvae']['kernel_size'])
+        self.decoder = Decoder(config['cvae']['decoder_input'], config['cvae']['decoder_depths'], config['cvae']['decoder_dims'], config['cvae']['kernel_size'], is_label=config['cvae']['is_label'])
 
+        self.add_noise = config['cvae']['add_noise']
+        self.noise_amount = config['cvae']['noise_amount']
 
     # @tf.function
     def sample_single(self, uv_template, uv_body_posed, uv_mask, eps=None, nr_of_samples=4):
@@ -122,33 +69,32 @@ class CVAE(tf.keras.Model):
         return x
 
 
-    def encode_conds(self, uv_template, uv_body_posed):
-        if uv_template.shape[1] == 3:
-            uv_template = tf.transpose(uv_template, [0,2,3,1])
-        if uv_body_posed.shape[1] == 3:
-            uv_body_posed = tf.transpose(uv_body_posed, [0,2,3,1])
-
-        z_static = self.static_enc(uv_template)
-        z_body = self.body_enc(uv_body_posed)
+    def encode_conds(self, uv_template, uv_body_posed, training):
+        z_static = self.static_enc(uv_template, training=training)
+        z_body = self.body_enc(uv_body_posed, training=training)
 
         z_cond = tf.concat([z_static, z_body], axis=-1)
-        z_cond = self.cond_encoder(z_cond)
+        z_cond = self.cond_encoder(z_cond, training=training)
 
         return z_cond
 
-    def sample(self, uv_template, uv_body_posed, uv_mask=None, noise=None):
+    def sample(self, uv_template, uv_body_posed, uv_mask, noise=None):
         if noise is None:
             eps = tf.random.normal(shape=(uv_template.shape[0], self.latent_dim))
         else:
             eps = noise
 
-        z_cond = self.encode_conds(uv_template, uv_body_posed)
+        z_cond = self.encode_conds(uv_template, uv_body_posed, training=False)
         z = tf.concat([eps, z_cond], axis=-1)
 
-        x = self.decoder(z)
+        x = self.decoder(z, training=False)
 
-        if uv_mask is not None:
-            x = x * uv_mask[:,:,:,tf.newaxis]
+        x = x * uv_mask[:,:,:,tf.newaxis]
+
+        background = tf.ones((uv_template.shape[0],512,256,1)) * (1 - uv_mask[:,:,:,tf.newaxis])
+        x = tf.concat((background, x), axis=-1)
+
+        x = tf.math.softmax(x)
 
         return x
 
@@ -171,32 +117,24 @@ class CVAE(tf.keras.Model):
         return eps * tf.exp(logvar * .5) + mean
 
 
-    def decode(self, z, apply_sigmoid=False):
-        logits = self.decoder(z)
-        if apply_sigmoid:
-            probs = tf.sigmoid(logits)
-            return probs
-        return logits
-
-
-    def call(self, inputs, training=None):
+    def call(self, inputs, training):
         uv_template, uv_body_posed, uv_curv, uv_mask = inputs
-        if uv_template.shape[1] == 3:
-            uv_template = tf.transpose(uv_template, [0,2,3,1])
-        if uv_body_posed.shape[1] == 3:
-            uv_body_posed = tf.transpose(uv_body_posed, [0,2,3,1])
-        if uv_curv.shape[1] == 3:
-            uv_curv = tf.transpose(uv_curv, [0,2,3,1])
 
-        z_cond = self.encode_conds(uv_template, uv_body_posed)
 
         x = tf.concat([uv_template, uv_body_posed, uv_curv], axis=-1)
         mean, logvar = self.encoder(x)
         z = self.reparameterize(mean, logvar)
 
+        z_cond = self.encode_conds(uv_template, uv_body_posed, training=training)
+
         z = tf.concat([z, z_cond], axis=-1)
-        x = self.decoder(z)
+        x = self.decoder(z, training=training)
 
         x = x * uv_mask[:,:,:,tf.newaxis]
 
-        return x, mean, logvar, z
+        background = tf.ones((4,512,256,1)) * (1 - uv_mask[:,:,:,tf.newaxis])
+        x = tf.concat((background, x), axis=-1)
+
+        x = tf.math.softmax(x)
+
+        return x, mean, logvar

models/cvae_model.py

@@ -9,6 +9,7 @@ import scipy as sp
 
 from custom_utils.custom_utils import load_model
 from models.base import AEEncoder, VariationalEncoder, Decoder
+from models.base import AEEncoderNew, VariationalEncoderNew, DecoderNew
 
 
 class CVAE(tf.keras.Model):
@@ -18,13 +19,14 @@ class CVAE(tf.keras.Model):
         self.latent_dim = config['cvae']['latent_dim']
         self.size = config['resolution']
 
-        self.static_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'])
-        if config['static_model'] != '':
-            load_model(self.static_enc, config['static_model'] + '_enc')
+        if not('only_curvatures' in config.keys() and config['only_curvatures'] == True):
+            self.static_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'])
+            if config['static_model'] != '':
+                load_model(self.static_enc, config['static_model'] + '_enc')
 
-        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['body_model'] != '':
-            load_model(self.body_enc, config['body_model'] + '_enc')
+            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['body_model'] != '':
+                load_model(self.body_enc, config['body_model'] + '_enc')
 
         if config['use_curvature']:
             self.curv_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'])
@@ -32,14 +34,16 @@ class CVAE(tf.keras.Model):
                 load_model(self.curv_enc, config['curvature_model'] + '_enc')
 
         if config['freeze_encoders']:
-            self.static_enc.trainable = False
-            self.body_enc.trainable = False
+            if not('only_curvatures' in config.keys() and config['only_curvatures'] == True):
+                self.static_enc.trainable = False
+                self.body_enc.trainable = False
             if config['use_curvature']:
                 self.curv_enc.trainable = False
 
         self.cond_encoder = tf.keras.Sequential([tf.keras.layers.Dense(512, activation='relu'),
                                                  tf.keras.layers.Dense(256, activation='relu'),
                                                  tf.keras.layers.Dense(128, activation='relu'),
+                                                 tf.keras.layers.Lambda(lambda x: x / 1000),
                                                  tf.keras.layers.Dense(config['latent_dim'], activation='tanh')])
 
         self.encoder = VariationalEncoder(config['cvae']['latent_dim'], config['cvae']['encoder_depths'], config['cvae']['encoder_dims'], kernel_size=config['cvae']['kernel_size'])
@@ -53,6 +57,7 @@ class CVAE(tf.keras.Model):
         self.add_noise = config['cvae']['add_noise']
         self.noise_amount = config['cvae']['noise_amount']
         self.use_curvature = config['use_curvature']
+        self.only_curvatures = config['only_curvatures'] if 'only_curvatures' in config.keys() and config['only_curvatures'] == True else None
         self.use_offsets_over_body = not config['use_skinning']
 
         if self.use_offsets_over_body:
@@ -60,14 +65,17 @@ class CVAE(tf.keras.Model):
         if self.avg_garment is not None:
             assert not self.use_offsets_over_body, 'When avg_garment is given, we must use skinning!'
 
-    @tf.function
-    def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask, noise=None):
+    # @tf.function
+    def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask, noise=None, add_noise=False):
         if noise is None:
             eps = tf.random.normal(shape=(uv_template.shape[0], self.latent_dim))
         else:
             eps = noise
 
         z_cond = self.encode_conds(uv_template, uv_body_posed, uv_curv, training=False)
+        if add_noise:
+            z_cond = z_cond + (tf.random.uniform(z_cond.shape, maxval=self.noise_amount * 2) - self.noise_amount)
+        
         z = tf.concat([eps, z_cond], axis=-1)
 
         x = self.decoder(z, training=False)
@@ -75,40 +83,187 @@ class CVAE(tf.keras.Model):
         if self.avg_garment is not None:
             x = x + self.avg_garment[tf.newaxis,:,:,:]
         
+        if self.use_offsets_over_body:
+            x = x + uv_body_posed
+
         x = x * uv_mask[:,:,:,tf.newaxis]
 
         return x
 
 
-    def reparameterize(self, mean, logvar, fix_noise=False):
-        if not fix_noise:
-            eps = tf.random.normal(shape=mean.shape)
+    def reparameterize(self, mean, logvar, training, fix_noise=False):
+        if training == True:
+            if not fix_noise:
+                eps = tf.random.normal(shape=mean.shape)
+            else:
+                eps = tf.random.normal(shape=(1, mean.shape[1]))
+                eps = tf.tile(eps, [mean.shape[0],1])
+            out = eps * tf.exp(logvar * .5) + mean
         else:
-            eps = tf.random.normal(shape=(1, mean.shape[1]))
-            eps = tf.tile(eps, [mean.shape[0],1])
-        return eps * tf.exp(logvar * .5) + mean
+            out = tf.exp(logvar * .5) + mean
+        return out
 
 
-    def encode_conds(self, uv_template, uv_body_posed, uv_curv, training=None):
-        z_body = self.body_enc(uv_body_posed, training=training)
-        z_static = self.static_enc(uv_template, training=training)
-        if self.use_curvature:
-            z_curv = self.curv_enc(uv_curv, training=training)
-            z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
+    def encode_conds(self, uv_template, uv_body_posed, uv_curv, training):
+        if self.only_curvatures:
+            z_cond = self.curv_enc(uv_curv, training=training)
         else:
-            z_cond = tf.concat([z_static, z_body], axis=-1)
+            z_body = self.body_enc(uv_body_posed, training=training)
+            z_static = self.static_enc(uv_template, training=training)
+            if self.use_curvature:
+                z_curv = self.curv_enc(uv_curv, training=training)
+                z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
+            else:
+                z_cond = tf.concat([z_static, z_body], axis=-1)
 
         z_cond = self.cond_encoder(z_cond, training=training)
 
         return z_cond
 
 
-    def call(self, inputs, training=None):
+    def call(self, inputs, training):
         uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask = inputs
-
-        x = tf.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], axis=-1)
+        if self.only_curvatures:
+            x = tf.concat([uv_unposed, uv_curv], axis=-1)
+        else:
+            x = tf.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], axis=-1)
+        
         mean, logvar = self.encoder(x, training=training)
-        z = self.reparameterize(mean, logvar)
+        z = self.reparameterize(mean, logvar, training)
+        z_cond = self.encode_conds(uv_template, uv_body_posed, uv_curv, training=training)
+
+        if self.add_noise and training == True:
+            z_cond = z_cond + (tf.random.uniform(z_cond.shape, maxval=self.noise_amount * 2) - self.noise_amount)
+
+
+        z = tf.concat([z, z_cond], axis=-1)
+        x = self.decoder(z, training=training)
+
+        if self.avg_garment is not None:
+            x = x + self.avg_garment[tf.newaxis,:,:,:]
+    
+        if self.use_offsets_over_body:
+            x = x + uv_body_posed
+
+        x = x * uv_mask[:,:,:,tf.newaxis]
+
+        return x, mean, logvar
+
+
+
+
+class CVAENew(tf.keras.Model):
+    
+    def __init__(self, config):
+        super(CVAENew, self).__init__()
+        self.size = config['resolution']
+
+        if not('only_curvatures' in config.keys() and config['only_curvatures'] == True):
+            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'])
+            if config['static_model'] != '':
+                load_model(self.static_enc, config['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'])
+            if config['body_model'] != '':
+                load_model(self.body_enc, config['body_model'] + '_enc')
+
+        if config['use_curvature']:
+            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'])
+            if config['curvature_model'] != '':
+                load_model(self.curv_enc, config['curvature_model'] + '_enc')
+
+        if config['freeze_encoders']:
+            if not('only_curvatures' in config.keys() and config['only_curvatures'] == True):
+                self.static_enc.trainable = False
+                self.body_enc.trainable = False
+            if config['use_curvature']:
+                self.curv_enc.trainable = False
+
+        self.encoder = VariationalEncoderNew(config['cvae']['encoder_depths'], config['cvae']['encoder_dims'], kernel_size=config['cvae']['kernel_size'])
+        self.decoder = DecoderNew(config['cvae']['decoder_depths'], config['cvae']['decoder_dims'], config['cvae']['kernel_size'])
+
+        self.reshape = tf.keras.layers.Reshape((8, 4, 64))
+
+        if config['avg_garment_path'] != '':
+            self.avg_garment = np.load(config['avg_garment_path'])
+        else:
+            self.avg_garment = None
+
+        self.add_noise = config['cvae']['add_noise']
+        self.noise_amount = config['cvae']['noise_amount']
+        self.use_curvature = config['use_curvature']
+        self.only_curvatures = config['only_curvatures'] if 'only_curvatures' in config.keys() and config['only_curvatures'] == True else None
+        self.use_offsets_over_body = not config['use_skinning']
+
+        if self.use_offsets_over_body:
+            assert self.avg_garment is None, 'When we are not using skinning, we can\'t use avg_garment!'
+        if self.avg_garment is not None:
+            assert not self.use_offsets_over_body, 'When avg_garment is given, we must use skinning!'
+
+    # @tf.function
+    def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask, noise=None, add_noise=False):
+        if noise is None:
+            eps = tf.random.normal(shape=(uv_template.shape[0], self.latent_dim))
+        else:
+            eps = noise
+
+        z_cond = self.encode_conds(uv_template, uv_body_posed, uv_curv, training=False)
+        if add_noise:
+            z_cond = z_cond + (tf.random.uniform(z_cond.shape, maxval=self.noise_amount * 2) - self.noise_amount)
+        
+        z = tf.concat([eps, z_cond], axis=-1)
+
+        x = self.decoder(z, training=False)
+        
+        if self.avg_garment is not None:
+            x = x + self.avg_garment[tf.newaxis,:,:,:]
+        
+        if self.use_offsets_over_body:
+            x = x + uv_body_posed
+
+        x = x * uv_mask[:,:,:,tf.newaxis]
+
+        return x
+
+
+    def reparameterize(self, mean, logvar, training, fix_noise=False):
+        if training == True:
+            if not fix_noise:
+                eps = tf.random.normal(shape=mean.shape)
+            else:
+                eps = tf.random.normal(shape=(1, mean.shape[1]))
+                eps = tf.tile(eps, [mean.shape[0],1])
+            out = eps * tf.exp(logvar * .5) + mean
+        else:
+            out = tf.exp(logvar * .5) + mean
+        return out
+
+
+    def encode_conds(self, uv_template, uv_body_posed, uv_curv, training):
+        if self.only_curvatures:
+            z_cond = self.curv_enc(uv_curv, training=training)
+        else:
+            z_body = self.body_enc(uv_body_posed, training=training)
+            z_static = self.static_enc(uv_template, training=training)
+            if self.use_curvature:
+                z_curv = self.curv_enc(uv_curv, training=training)
+                z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
+            else:
+                z_cond = tf.concat([z_static, z_body], axis=-1)
+
+        return z_cond
+
+
+    def call(self, inputs, training):
+        uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask = inputs
+        if self.only_curvatures:
+            x = tf.concat([uv_unposed, uv_curv], axis=-1)
+        else:
+            x = tf.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], axis=-1)
+        
+        mean, logvar = self.encoder(x, training=training)
+        z = self.reparameterize(mean, logvar, training)
+        z = self.reshape(z)
         z_cond = self.encode_conds(uv_template, uv_body_posed, uv_curv, training=training)
 
         if self.add_noise and training == True:

models/cvae_model_overfit_jumpsuit.py

@@ -0,0 +1,142 @@
+import os
+import time
+import random
+import tensorflow as tf
+import numpy as np
+import scipy as sp
+# import plotly.graph_objects as go
+# import plotly.express as ex
+
+from custom_utils.custom_utils import load_model
+from models.base import AEEncoder, VariationalEncoder, Decoder
+
+
+class CVAE(tf.keras.Model):
+    
+    def __init__(self, config):
+        super(CVAE, self).__init__()
+        self.latent_dim = config['cvae']['latent_dim']
+        self.size = config['resolution']
+
+        if not('only_curvatures' in config.keys() and config['only_curvatures'] == True):
+            self.static_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'])
+            if config['static_model'] != '':
+                load_model(self.static_enc, config['static_model'] + '_enc')
+
+            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['body_model'] != '':
+                load_model(self.body_enc, config['body_model'] + '_enc')
+
+        if config['use_curvature']:
+            self.curv_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'])
+            if config['curvature_model'] != '':
+                load_model(self.curv_enc, config['curvature_model'] + '_enc')
+
+        if config['freeze_encoders']:
+            if not('only_curvatures' in config.keys() and config['only_curvatures'] == True):
+                self.static_enc.trainable = False
+                self.body_enc.trainable = False
+            if config['use_curvature']:
+                self.curv_enc.trainable = False
+
+        self.cond_encoder = tf.keras.Sequential([tf.keras.layers.Dense(512, activation='relu'),
+                                                 tf.keras.layers.Dense(256, activation='relu'),
+                                                 tf.keras.layers.Dense(128, activation='relu'),
+                                                 tf.keras.layers.Dense(config['latent_dim'], activation='tanh')])
+
+        self.encoder = VariationalEncoder(config['cvae']['latent_dim'], config['cvae']['encoder_depths'], config['cvae']['encoder_dims'], kernel_size=config['cvae']['kernel_size'])
+        self.decoder = Decoder(config['cvae']['decoder_input'], config['cvae']['decoder_depths'], config['cvae']['decoder_dims'], config['cvae']['kernel_size'])
+
+        if config['avg_garment_path'] != '':
+            self.avg_garment = np.load(config['avg_garment_path'])
+        else:
+            self.avg_garment = None
+
+        self.add_noise = config['cvae']['add_noise']
+        self.noise_amount = config['cvae']['noise_amount']
+        self.use_curvature = config['use_curvature']
+        self.only_curvatures = config['only_curvatures'] if 'only_curvatures' in config.keys() and config['only_curvatures'] == True else None
+        self.use_offsets_over_body = not config['use_skinning']
+
+        if self.use_offsets_over_body:
+            assert self.avg_garment is None, 'When we are not using skinning, we can\'t use avg_garment!'
+        if self.avg_garment is not None:
+            assert not self.use_offsets_over_body, 'When avg_garment is given, we must use skinning!'
+
+    # @tf.function
+    def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask, noise=None):
+        if noise is None:
+            eps = tf.random.normal(shape=(uv_template.shape[0], self.latent_dim))
+        else:
+            eps = noise
+
+        z_cond = self.encode_conds(uv_template, uv_body_posed, uv_curv, training=False)
+        z = tf.concat([eps, z_cond], axis=-1)
+
+        x = self.decoder(z, training=False)
+        
+        if self.avg_garment is not None:
+            x = x + self.avg_garment[tf.newaxis,:,:,:]
+        
+        if self.use_offsets_over_body:
+            x = x + uv_body_posed
+
+        x = x * uv_mask[:,:,:,tf.newaxis]
+
+        return x
+
+
+    def reparameterize(self, mean, logvar, training, fix_noise=False):
+        if training == True:
+            if not fix_noise:
+                eps = tf.random.normal(shape=mean.shape)
+            else:
+                eps = tf.random.normal(shape=(1, mean.shape[1]))
+                eps = tf.tile(eps, [mean.shape[0],1])
+            out = eps * tf.exp(logvar * .5) + mean
+        else:
+            out = tf.exp(logvar * .5) + mean
+        return out
+
+
+    def encode_conds(self, uv_template, uv_body_posed, uv_curv, training):
+        if self.only_curvatures:
+            z_cond = self.curv_enc(uv_curv, training=training)
+        else:
+            z_body = self.body_enc(uv_body_posed, training=training)
+            z_static = self.static_enc(uv_template, training=training)
+            if self.use_curvature:
+                z_curv = self.curv_enc(uv_curv, training=training)
+                z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
+            else:
+                z_cond = tf.concat([z_static, z_body], axis=-1)
+
+        z_cond = self.cond_encoder(z_cond, training=training)
+
+        return z_cond
+
+
+    def call(self, inputs, training):
+        uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask = inputs
+
+        x = tf.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], axis=-1)
+        mean, logvar = self.encoder(x, training=training)
+        z = self.reparameterize(mean, logvar, training)
+        z_cond = self.encode_conds(uv_template, uv_body_posed, uv_curv, training=training)
+
+        if self.add_noise and training == True:
+            z_cond = z_cond + (tf.random.uniform(z_cond.shape, maxval=self.noise_amount * 2) - self.noise_amount)
+
+
+        z = tf.concat([z, z_cond], axis=-1)
+        x = self.decoder(z, training=training)
+
+        if self.avg_garment is not None:
+            x = x + self.avg_garment[tf.newaxis,:,:,:]
+    
+        if self.use_offsets_over_body:
+            x = x + uv_body_posed
+
+        x = x * uv_mask[:,:,:,tf.newaxis]
+
+        return x, mean, logvar

models/cvae_pyramid_model.py

@@ -60,13 +60,13 @@ class Pyramid(tf.keras.Model):
         else:
             self.avg_garment = None
 
-    def enhance(self, low, mask, training=None):
+    def enhance(self, low, mask, training):
         x = self.upscale_enhancer(low, training=training)
         x = x * (1 - mask[:,:,:,tf.newaxis])
         low = x + low
         return low
 
-    def call(self, inputs, training=None):
+    def call(self, inputs, training):
         uv_unposed, uv_static, uv_body_posed, uv_curv, uv_mask, uv_unposed_low, uv_static_low, uv_body_posed_low, uv_curv_low, uv_mask_low = inputs
 
         if self.use_full_res_encoders:

models/pyramid_level.py

@@ -8,7 +8,7 @@ import scipy as sp
 # import plotly.express as ex
 
 from custom_utils.custom_utils import load_model
-from models.base import AEEncoder, VariationalEncoder, Decoder
+from models.base import AEEncoder, VariationalEncoder, Decoder, VariationalEncoderNew, DecoderNew
 from models.upscaler import UpscalerBasic, UpscalerBasicConv, UpscalerBasicNew
 
 
@@ -24,7 +24,8 @@ class PyramidLevel(tf.keras.Model):
 
         self.upscale_enhancer = None
         if config['upscale_enhancer'] != '':
-            self.upscale_enhancer = UpscalerBasicConv(config['resolution'])
+            k_size = config['upscale_enhancer_kernel_size'] if 'upscale_enhancer_kernel_size' in config.keys() else ((3 if config['resolution'] <= 64 else 5))
+            self.upscale_enhancer = UpscalerBasicConv(config['resolution'], kernel_size=k_size)
             load_model(self.upscale_enhancer, config['upscale_enhancer'])
             self.upscale_enhancer.trainable = not config['upscale_enhancer_freeze']
 
@@ -57,13 +58,17 @@ class PyramidLevel(tf.keras.Model):
         return x
 
 
-    def reparameterize(self, mean, logvar, fix_noise=False):
-        if not fix_noise:
-            eps = tf.random.normal(shape=mean.shape)
+    def reparameterize(self, mean, logvar, training, fix_noise=False):
+        if training == True:
+            if not fix_noise:
+                eps = tf.random.normal(shape=mean.shape)
+            else:
+                eps = tf.random.normal(shape=(1, mean.shape[1]))
+                eps = tf.tile(eps, [mean.shape[0],1])
+            out = eps * tf.exp(logvar * .5) + mean
         else:
-            eps = tf.random.normal(shape=(1, mean.shape[1]))
-            eps = tf.tile(eps, [mean.shape[0],1])
-        return eps * tf.exp(logvar * .5) + mean
+            out = tf.exp(logvar * .5) + mean
+        return out
 
 
     def call(self, inputs, z_cond, training=None):
@@ -71,7 +76,7 @@ class PyramidLevel(tf.keras.Model):
 
         x = tf.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], axis=-1)
         mean, logvar = self.encoder(x, training=training)
-        z = self.reparameterize(mean, logvar)
+        z = self.reparameterize(mean, logvar, training)
 
         z = tf.concat([z, z_cond], axis=-1)
         x = self.decoder(z, training=training)
@@ -89,3 +94,88 @@ class PyramidLevel(tf.keras.Model):
 
         return x, mean, logvar
 
+
+
+class PyramidLevelNew(tf.keras.Model):
+    
+    def __init__(self, config):
+        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.upscale_enhancer = None
+        if config['upscale_enhancer'] != '':
+            k_size = config['upscale_enhancer_kernel_size'] if 'upscale_enhancer_kernel_size' in config.keys() else ((3 if config['resolution'] <= 64 else 5))
+            self.upscale_enhancer = UpscalerBasicConv(config['resolution'], kernel_size=k_size)
+            load_model(self.upscale_enhancer, config['upscale_enhancer'])
+            self.upscale_enhancer.trainable = not config['upscale_enhancer_freeze']
+
+        if config['avg_garment_path'] != '':
+            self.avg_garment = np.load(config['avg_garment_path'])
+        else:
+            self.avg_garment = None
+
+        self.reshape = tf.keras.layers.Reshape((8, 4, 64))
+
+
+    def enhance(self, low, mask, training=None):
+        x = self.upscale_enhancer([low, mask], training=training)
+        return x
+
+
+    def sample(self, z_cond, uv_mask, noise=None):
+        if noise is None:
+            eps = tf.random.normal(shape=(uv_mask.shape[0], self.latent_dim))
+        else:
+            eps = noise
+
+        z = tf.concat([eps, z_cond], axis=-1)
+
+        x = self.decoder(z)
+        
+        if self.avg_garment is not None:
+            x = x + self.avg_garment[tf.newaxis,:,:,:]
+        
+        x = x * uv_mask[:,:,:,tf.newaxis]
+
+        return x
+
+
+    def reparameterize(self, mean, logvar, training, fix_noise=False):
+        if training == True:
+            if not fix_noise:
+                eps = tf.random.normal(shape=mean.shape)
+            else:
+                eps = tf.random.normal(shape=(1, mean.shape[1]))
+                eps = tf.tile(eps, [mean.shape[0],1])
+            out = eps * tf.exp(logvar * .5) + mean
+        else:
+            out = tf.exp(logvar * .5) + mean
+        return out
+
+
+    def call(self, inputs, z_cond, training=None):
+        uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask, uv_mask_high = inputs
+
+        x = tf.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], axis=-1)
+        mean, logvar = self.encoder(x, training=training)
+        z = self.reparameterize(mean, logvar, training)
+        z = self.reshape(z)
+
+        z = tf.concat([z, z_cond], axis=-1)
+        x = self.decoder(z, training=training)
+
+        if self.avg_garment is not None:
+            x = x + self.avg_garment[tf.newaxis,:,:,:]
+
+        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)
+        
+        # x = x * uv_mask_high[:,:,:,tf.newaxis]
+
+        return x, mean, logvar

models/pyramid_model.py

@@ -49,7 +49,7 @@ class Pyramid(tf.keras.Model):
             load_model(self.upscale_enhancer, config['pyramid']['final_upscale_enhancer'])
             self.upscale_enhancer.trainable = False
 
-    def enhance(self, low, mask, training=None):
+    def enhance(self, low, mask, training):
         x = self.upscale_enhancer([low, mask], training=training)
         return x
 
@@ -63,7 +63,7 @@ class Pyramid(tf.keras.Model):
 
         return z_cond
 
-    def call(self, inputs, training=None):
+    def call(self, inputs, training):
         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)
@@ -136,7 +136,7 @@ class Pyramid(tf.keras.Model):
 #             load_model(self.upscale_enhancer, config['pyramid']['final_upscale_enhancer'])
 #             self.upscale_enhancer.trainable = False
 
-#     def enhance(self, low, mask, training=None):
+#     def enhance(self, low, mask, training):
 #         x = self.upscale_enhancer([low, mask], training=training)
 #         return x
 
@@ -150,7 +150,7 @@ class Pyramid(tf.keras.Model):
 
 #         return z_cond
 
-#     def call(self, inputs, training=None):
+#     def call(self, inputs, training):
 #         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)

models/pyramid_model_optimized.py

@@ -7,9 +7,8 @@ import numpy as np
 # import plotly.express as ex
 
 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.base import AEEncoder, AEEncoderNew
+from models.pyramid_level import PyramidLevel, PyramidLevelNew
 from models.upscaler import UpscalerBasic, UpscalerBasicConv, UpscalerBasicNew
 
 
@@ -42,57 +41,65 @@ 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.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
-        # self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
-        # self.levels[-3].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
-
+        # if config['is_cape']:
+        #     if len(config['pyramid']['resolutions']) == 2:
+        #         self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
+        #     if len(config['pyramid']['resolutions']) == 3:
+        #         self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
+        #         self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
+        #     if len(config['pyramid']['resolutions']) == 4:
+        #         self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
+        #         self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
+        #         self.levels[-3].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
+        # else:
         if len(config['pyramid']['resolutions']) == 2:
-            self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 10)])
-        if len(config['pyramid']['resolutions']) == 3:
             self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
-            self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 10)])
-        if len(config['pyramid']['resolutions']) == 4:
+        if len(config['pyramid']['resolutions']) == 3:
             self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 1000)])
             self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
-            self.levels[-3].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 10)])
-        if len(config['pyramid']['resolutions']) == 5:
+        if len(config['pyramid']['resolutions']) == 4:
             self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 10000)])
             self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 1000)])
             self.levels[-3].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
-            self.levels[-4].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 10)])
+        if len(config['pyramid']['resolutions']) == 5:
+            self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100000)])
+            self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 10000)])
+            self.levels[-3].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 1000)])
+            self.levels[-4].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
 
         self.add_noise = config['pyramid']['add_noise']
+        self.noise_amount = config['pyramid']['noise_amount']
         self.use_offsets = config['use_offsets']
  
         self.upscale_enhancer = None
         if config['pyramid']['final_upscale_enhancer'] != '':
-            self.upscale_enhancer = UpscalerBasicConv(low_res=config['pyramid']['resolutions'][-1])
+            k_size = config['pyramid']['final_upscale_enhancer_kernel_size'] if 'final_upscale_enhancer_kernel_size' in config['pyramid'].keys() else ((3 if config['pyramid']['resolutions'][-1] <= 64 else 5))
+            self.upscale_enhancer = UpscalerBasicConv(low_res=config['pyramid']['resolutions'][-1], kernel_size=k_size)
             load_model(self.upscale_enhancer, config['pyramid']['final_upscale_enhancer'])
             self.upscale_enhancer.trainable = False
 
-    def enhance(self, low, mask, training=None):
+    def enhance(self, low, mask, training):
         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)
-        z_body = self.body_enc(uv_body_posed)
-        z_curv = self.curv_enc(uv_curv)
+    def encode_conds(self, uv_template, uv_body_posed, uv_curv, training):
+        z_static = self.static_enc(uv_template, training)
+        z_body = self.body_enc(uv_body_posed, training)
+        z_curv = self.curv_enc(uv_curv, training)
 
         z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
-        z_cond = self.cond_encoder(z_cond)
+        z_cond = self.cond_encoder(z_cond, training=training)
 
         return z_cond
 
-    def call(self, inputs, training=None):
+    def call(self, inputs, training):
         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)
+        z_cond = self.encode_conds(uv_static, uv_body, uv_curv, training)
 
-        if self.add_noise:
-            z_cond = z_cond + tf.random.uniform(z_cond.shape, maxval=10)
+        if self.add_noise == True and training == True:
+            z_cond = z_cond + (tf.random.uniform(z_cond.shape, maxval=self.noise_amount * 2) - self.noise_amount)
 
         out_levels = []
         pred, mean, logvar = self.levels[0](levels_input[0], z_cond, training=training)
@@ -184,7 +191,7 @@ class Pyramid(tf.keras.Model):
 #             load_model(self.upscale_enhancer, config['pyramid']['final_upscale_enhancer'])
 #             self.upscale_enhancer.trainable = False
 
-#     def enhance(self, low, mask, training=None):
+#     def enhance(self, low, mask, training):
 #         x = self.upscale_enhancer([low, mask], training=training)
 #         return x
 
@@ -198,7 +205,7 @@ class Pyramid(tf.keras.Model):
 
 #         return z_cond
 
-#     def call(self, inputs, training=None):
+#     def call(self, inputs, training):
 #         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)
@@ -241,4 +248,118 @@ class Pyramid(tf.keras.Model):
 #             [pred_0, mean_0, logvar_0]
 #         ]
 
-#         return out_levels
+#         return out_levels
+
+
+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'])
+        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'])
+        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'])
+        if config['pyramid']['curvature_model'] != '':
+            load_model(self.curv_enc, config['pyramid']['curvature_model'] + '_enc')
+
+        if config['pyramid']['freeze_encoders']:
+            self.static_enc.trainable = False
+            self.body_enc.trainable = False
+            self.curv_enc.trainable = False
+
+        self.levels = [PyramidLevelNew({ **config['pyramid']['levels'][i] }) for i in range(len(config['pyramid']['levels']))]
+
+
+        # if len(config['pyramid']['resolutions']) == 2:
+        #     self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
+        # if len(config['pyramid']['resolutions']) == 3:
+        #     self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 1000)])
+        #     self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
+        # if len(config['pyramid']['resolutions']) == 4:
+        #     self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 10000)])
+        #     self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 1000)])
+        #     self.levels[-3].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
+        # if len(config['pyramid']['resolutions']) == 5:
+        #     self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100000)])
+        #     self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 10000)])
+        #     self.levels[-3].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 1000)])
+        #     self.levels[-4].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
+
+        self.add_noise = config['pyramid']['add_noise']
+        self.noise_amount = config['pyramid']['noise_amount']
+        self.use_offsets = config['use_offsets']
+ 
+        self.upscale_enhancer = None
+        if config['pyramid']['final_upscale_enhancer'] != '':
+            k_size = config['pyramid']['final_upscale_enhancer_kernel_size'] if 'final_upscale_enhancer_kernel_size' in config['pyramid'].keys() else ((3 if config['pyramid']['resolutions'][-1] <= 64 else 5))
+            self.upscale_enhancer = UpscalerBasicConv(low_res=config['pyramid']['resolutions'][-1], kernel_size=k_size)
+            load_model(self.upscale_enhancer, config['pyramid']['final_upscale_enhancer'])
+            self.upscale_enhancer.trainable = False
+
+    def enhance(self, low, mask, training):
+        x = self.upscale_enhancer([low, mask], training=training)
+        return x
+
+    def encode_conds(self, uv_template, uv_body_posed, uv_curv, training):
+        z_static = self.static_enc(uv_template, training)
+        z_body = self.body_enc(uv_body_posed, training)
+        z_curv = self.curv_enc(uv_curv, training)
+
+        z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
+
+        return z_cond
+
+    def call(self, inputs, training):
+        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, training)
+
+        if self.add_noise == True and training == True:
+            z_cond = z_cond + (tf.random.uniform(z_cond.shape, maxval=self.noise_amount * 2) - self.noise_amount)
+
+        out_levels = []
+        pred, mean, logvar = self.levels[0](levels_input[0], z_cond, training=training)
+        if self.use_offsets:
+            pred = pred + levels_input[0][2] # uv_body_posed_32
+        if 1 != len(self.levels):
+            pred = self.levels[0].enhance(pred, levels_input[0][-2], training=training)  # levels_input[i][-2]=uv_mask_res
+            pred = tf.image.resize(pred, (self.levels[1].size, self.levels[1].size//2), tf.image.ResizeMethod.BILINEAR)
+            out_levels.append([pred, mean, logvar])
+            pred = pred * levels_input[1][-2][:,:,:,tf.newaxis]
+        else:
+            # pred = self.enhance(pred, levels_input[0][-2], training=training)
+            pred = self.levels[0].enhance(pred, levels_input[0][-2], training=training)  # levels_input[i][-2]=uv_mask_res
+            pred = tf.image.resize(pred, (512, 256), tf.image.ResizeMethod.BILINEAR)
+            pred = pred * levels_input[-1][-1][:,:,:,tf.newaxis] # full res uv mask
+            out_levels.append([pred, mean, logvar])
+
+        for i in range(1, len(self.levels)):
+            input = levels_input[i]
+            input[0] = input[0] - pred
+            input_min, input_max = tf.reduce_min(input[0], axis=[1,2,3]), tf.reduce_max(input[0], axis=[1,2,3])
+            input[0] = (input[0] - input_min[:,None,None,None]) / (input_max - input_min)[:,None,None,None]
+            input[0] = input[0] * input[-2][:,:,:,None] # uv_mask_res
+            pred_level, mean, logvar = self.levels[i](input, z_cond, training=training)
+            pred_level = ( (pred_level * (input_max - input_min)[:,None,None,None]) + input_min[:,None,None,None] ) * input[-2][:,:,:,None] # uv_mask_res
+            pred = pred + pred_level
+            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.BILINEAR)
+                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 = self.levels[-1].enhance(pred, levels_input[-1][-2], training=training)
+                pred = tf.image.resize(pred, (512, 256), tf.image.ResizeMethod.BILINEAR)
+                pred = pred * levels_input[-1][-1][:,:,:,tf.newaxis] # full res uv mask
+                out_levels.append([pred, mean, logvar])
+
+        return out_levels

models/upscaler.py

@@ -12,7 +12,7 @@ class Bottleneck(tf.keras.Model):
                        layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
             )
 
-    def call(self, inputs, training=None):
+    def call(self, inputs, training):
         x = inputs
 
         x = self.bottleneck(x, training=training)
@@ -35,7 +35,7 @@ class Upscaler(tf.keras.Model):
 
         self.decoder = ConvNeXtReverseWithSkip(depths=depths, dims=reverse_dims, kernel_size=3)
 
-    def call(self, inputs, training=None):
+    def call(self, inputs, training):
         img, mask = inputs
 
         img = tf.transpose(img, [0,2,3,1])
@@ -74,7 +74,7 @@ class UpscalerBasic(tf.keras.Model):
 
         self.reshape = tf.keras.layers.Reshape((low_res, (low_res//2), 3))
 
-    def call(self, inputs, training=None):
+    def call(self, inputs, training):
         img, mask = inputs
 
         x = self.flatten(img)
@@ -110,7 +110,7 @@ class UpscalerBasicNew(tf.keras.Model):
 
         self.reshape = tf.keras.layers.Reshape((low_res, (low_res//2), 3), name='upscale_enhancer_reshape')
 
-    def call(self, inputs, training=None):
+    def call(self, inputs, training):
         img, mask = inputs
 
         x = self.flatten(img)
@@ -143,7 +143,7 @@ class UpscalerBasicConv(tf.keras.Model):
         ])
 
 
-    def call(self, inputs, training=None):
+    def call(self, inputs, training):
         img, mask = inputs
 
         x = self.seq(img, training=training)
@@ -160,7 +160,7 @@ class UpscalerBasicConv(tf.keras.Model):
 
 #         self.decoder = ConvNeXtReverse(depths=list(reversed([9, 9, 9])), dims=list(reversed([bottleneck_size, 128, 32, 16])))
 
-#     def call(self, inputs, training=None):
+#     def call(self, inputs, training):
 #         x, mask = inputs
 
 #         x = tf.transpose(x, [0,2,3,1])