large refactor

models/ae_curv_model.py

@@ -1,63 +1,73 @@
 import tensorflow as tf
 from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverseLabels, Block
+from models.base import AEEncoder, Bottleneck, Decoder
 
 
 
-class Bottleneck(tf.keras.Model):
+# class Bottleneck(tf.keras.Model):
     
-    def __init__(self, dim, latent_dim):
-        super(Bottleneck, self).__init__()
+#     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)
+#         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
+#     def call(self, inputs, training=None):
+#         x = inputs
 
-        x = self.bottleneck(x)
-        x = self.flatten(x)
-        x = self.dense(x)
+#         x = self.bottleneck(x)
+#         x = self.flatten(x)
+#         x = self.dense(x)
         
-        return x
+#         return x
 
 
-class Decoder(tf.keras.Model):
+# class Decoder(tf.keras.Model):
     
-    def __init__(self, bottleneck_size, latent_dim):
-        super(Decoder, self).__init__()
+#     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)
+#         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
+#     def call(self, inputs, training=None):
+#         x = inputs
 
-        x = self.dense(x)
-        x = self.reshape(x)
-        x = self.decoder(x)
+#         x = self.dense(x)
+#         x = self.reshape(x)
+#         x = self.decoder(x)
 
-        return x
+#         return x
 
         
 class AutoEncoder(tf.keras.Model):
     
-    def __init__(self):
+    def __init__(self, size=512):
         super(AutoEncoder, self).__init__()
 
+        if size == 512:
+            bottleneck_size = 64
+            latent_dim = 128
+            kernel_size = 7
+            encoder_depths = [3, 3, 9, 3, 3]
+            encoder_dims = [32, 64, 128, 256, bottleneck_size]
+            decoder_input = (16, 8, bottleneck_size)
+            decoder_depths = [3, 3, 9, 3, 3]
+            decoder_dims = [bottleneck_size, 256, 128, 64, 16]
+        elif size == 128:
+            raise NotImplementedError
+        elif size == 32:
+            raise NotImplementedError
+        else:
+            raise NotImplementedError
 
-        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)
-
+        self.encoder : AEEncoder = AEEncoder(latent_dim, encoder_depths, encoder_dims, bottleneck_size, kernel_size)
+        self.decoder = Decoder(decoder_input, decoder_depths, decoder_dims, kernel_size, is_label=True)
 
 
     def call(self, inputs, training=None):
@@ -65,117 +75,113 @@ class AutoEncoder(tf.keras.Model):
 
         # 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]
+        x = x * uv_mask[:,:,:,tf.newaxis]
 
         return x
 
 
 
-class Decoder32(tf.keras.Model):
+# class Decoder32(tf.keras.Model):
     
-    def __init__(self, bottleneck_size, latent_dim):
-        super(Decoder32, self).__init__()
+#     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)
+#         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
+#     def call(self, inputs, training=None):
+#         x = inputs
 
-        x = self.dense(x)
-        x = self.reshape(x)
-        x = self.decoder(x)
+#         x = self.dense(x)
+#         x = self.reshape(x)
+#         x = self.decoder(x)
 
-        return x
+#         return x
 
         
-class AutoEncoder32(tf.keras.Model):
+# class AutoEncoder32(tf.keras.Model):
     
-    def __init__(self):
-        super(AutoEncoder32, self).__init__()
+#     def __init__(self):
+#         super(AutoEncoder32, self).__init__()
 
-        bottleneck_size = 512
-        latent_dim = 128
+#         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)
+#         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
+#     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 = tf.transpose(x, [0,2,3,1])
+#         x = self.encoder(x, training=training)
 
-        x = self.bottleneck(x, training=training)
+#         x = self.bottleneck(x, training=training)
 
-        x = self.decoder(x, training=training)
-        # x = tf.transpose(x, [0,3,1,2])
+#         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]
+#         # x = x * uv_mask[:,:, tf.newaxis]
+#         # x = x * uv_mask[:,tf.newaxis,:,:]
+#         # x = x * uv_mask[:,:,:,tf.newaxis]
 
-        return x
+#         return x
 
 
 
-class Decoder128(tf.keras.Model):
+# class Decoder128(tf.keras.Model):
     
-    def __init__(self, bottleneck_size, latent_dim):
-        super(Decoder128, self).__init__()
+#     def __init__(self, bottleneck_size, latent_dim):
+#         super(Decoder128, 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([16, 16, 9, 9])), dims=list(reversed([bottleneck_size, 256, 128, 64])), nr_of_classes=7)
+#         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([16, 16, 9, 9])), dims=list(reversed([bottleneck_size, 256, 128, 64])), nr_of_classes=7)
 
-    def call(self, inputs, training=None):
-        x = inputs
+#     def call(self, inputs, training=None):
+#         x = inputs
 
-        x = self.dense(x)
-        x = self.reshape(x)
-        x = self.decoder(x)
+#         x = self.dense(x)
+#         x = self.reshape(x)
+#         x = self.decoder(x)
 
-        return x
+#         return x
 
         
-class AutoEncoder128(tf.keras.Model):
+# class AutoEncoder128(tf.keras.Model):
     
-    def __init__(self):
-        super(AutoEncoder128, self).__init__()
+#     def __init__(self):
+#         super(AutoEncoder128, self).__init__()
 
-        bottleneck_size = 512
-        latent_dim = 128
+#         bottleneck_size = 512
+#         latent_dim = 128
 
-        self.encoder = ConvNeXt(depths=[16, 16, 9, 9], dims=[64, 128, 256, bottleneck_size], include_top=False)
-        self.bottleneck = Bottleneck(bottleneck_size, latent_dim)
-        self.decoder = Decoder128(bottleneck_size, latent_dim)
+#         self.encoder = ConvNeXt(depths=[16, 16, 9, 9], dims=[64, 128, 256, bottleneck_size], include_top=False)
+#         self.bottleneck = Bottleneck(bottleneck_size, latent_dim)
+#         self.decoder = Decoder128(bottleneck_size, latent_dim)
 
 
-    def call(self, inputs, training=None):
-        x, uv_mask = inputs
+#     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 = tf.transpose(x, [0,2,3,1])
+#         x = self.encoder(x, training=training)
 
-        x = self.bottleneck(x, training=training)
+#         x = self.bottleneck(x, training=training)
 
-        x = self.decoder(x, training=training)
-        # x = tf.transpose(x, [0,3,1,2])
+#         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]
+#         # x = x * uv_mask[:,:, tf.newaxis]
+#         # x = x * uv_mask[:,tf.newaxis,:,:]
+#         # x = x * uv_mask[:,:,:,tf.newaxis]
 
-        return x
+#         return x