commit all

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]