separate sub models from their training (wip)

models/ae_curv_model.py

@@ -0,0 +1,130 @@
+import tensorflow as tf
+from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverseLabels, Block
+
+
+
+class Bottleneck(tf.keras.Model):
+    
+    def __init__(self, dim, latent_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)]
+            )
+        self.flatten = tf.keras.layers.Flatten()
+        self.dense = tf.keras.layers.Dense(latent_dim)
+
+    def call(self, inputs, training=None):
+        x = inputs
+
+        x = self.bottleneck(x)
+        x = self.flatten(x)
+        x = self.dense(x)
+        
+        return x
+
+
+class Decoder(tf.keras.Model):
+    
+    def __init__(self, bottleneck_size, latent_dim):
+        super(Decoder, self).__init__()
+
+        self.dense = tf.keras.layers.Dense(16 * 8 * bottleneck_size)
+        self.reshape = tf.keras.layers.Reshape((16, 8, bottleneck_size))
+        self.decoder = ConvNeXtReverseLabels(depths=list(reversed([3, 3, 9, 3])), dims=list(reversed([bottleneck_size, 128, 64, 16])), nr_of_classes=7)
+
+    def call(self, inputs, training=None):
+        x = inputs
+
+        x = self.dense(x)
+        x = self.reshape(x)
+        x = self.decoder(x)
+
+        return x
+
+        
+class AutoEncoder(tf.keras.Model):
+    
+    def __init__(self):
+        super(AutoEncoder, self).__init__()
+
+
+        bottleneck_size = 32
+        latent_dim = 128
+        self.decoder = Decoder(bottleneck_size, latent_dim)
+
+        self.bottleneck = Bottleneck(bottleneck_size, latent_dim)
+
+        self.encoder = ConvNeXt(depths=[3, 3, 9, 3], dims=[32, 64, 128, bottleneck_size], include_top=False)
+
+
+
+    def call(self, inputs, training=None):
+        x, uv_mask = inputs
+
+        # x = tf.transpose(x, [0,2,3,1])
+        x = self.encoder(x, training=training)
+
+
+        x = self.bottleneck(x, training=training)
+
+        x = self.decoder(x, training=training)
+        # x = tf.transpose(x, [0,3,1,2])
+
+        # x = x * uv_mask[:,:, tf.newaxis]
+        # x = x * uv_mask[:,tf.newaxis,:,:]
+        x = x * uv_mask[:,:,:,tf.newaxis]
+
+        return x
+
+
+
+class Decoder32(tf.keras.Model):
+    
+    def __init__(self, bottleneck_size, latent_dim):
+        super(Decoder32, self).__init__()
+
+        self.dense = tf.keras.layers.Dense(4 * 2 * bottleneck_size)
+        self.reshape = tf.keras.layers.Reshape((4, 2, bottleneck_size))
+        self.decoder = ConvNeXtReverseLabels(depths=list(reversed([27, 27])), dims=list(reversed([bottleneck_size, 256])), nr_of_classes=7)
+
+    def call(self, inputs, training=None):
+        x = inputs
+
+        x = self.dense(x)
+        x = self.reshape(x)
+        x = self.decoder(x)
+
+        return x
+
+        
+class AutoEncoder32(tf.keras.Model):
+    
+    def __init__(self):
+        super(AutoEncoder32, self).__init__()
+
+        bottleneck_size = 512
+        latent_dim = 128
+
+        self.encoder = ConvNeXt(depths=[27, 27], dims=[256, bottleneck_size], include_top=False)
+        self.bottleneck = Bottleneck(bottleneck_size, latent_dim)
+        self.decoder = Decoder32(bottleneck_size, latent_dim)
+
+
+    def call(self, inputs, training=None):
+        x, uv_mask = inputs
+
+        # x = tf.transpose(x, [0,2,3,1])
+        x = self.encoder(x, training=training)
+
+        x = self.bottleneck(x, training=training)
+
+        x = self.decoder(x, training=training)
+        # x = tf.transpose(x, [0,3,1,2])
+
+        # x = x * uv_mask[:,:, tf.newaxis]
+        # x = x * uv_mask[:,tf.newaxis,:,:]
+        # x = x * uv_mask[:,:,:,tf.newaxis]
+
+        return x