apply for body and curv

models/base.py

@@ -152,7 +152,6 @@ 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.encoder = resnet_18()
         self.act = tf.keras.layers.Activation(activation)
 
@@ -180,10 +179,9 @@ class DecoderNew(tf.keras.Model):
         self.is_label = is_label
 
         if not is_label:
-            # self.decoder = ConvNeXtReverse(depths=depths, dims=dims, kernel_size=kernel_size)
             self.decoder = resnet_18(reversed=True)
         else:
-            self.decoder = ConvNeXtReverseLabels(depths=depths, dims=dims, nr_of_classes=7)
+            self.decoder = resnet_18(reversed=True, labels=6)
             
 
     def call(self, inputs, training):

models/resnet.py

@@ -116,9 +116,10 @@ def make_bottleneck_layer(filter_num, blocks, stride=1):
 
 
 class ResNetTypeI(tf.keras.Model):
-    def __init__(self, layer_params, reversed=False):
+    def __init__(self, layer_params, reversed=False, labels=None):
         super(ResNetTypeI, self).__init__()
         self.reversed = reversed
+        self.labels = labels
 
         if not reversed:
             self.conv1 = tf.keras.layers.Conv2D(filters=64,
@@ -130,7 +131,7 @@ class ResNetTypeI(tf.keras.Model):
                                                 strides=2,
                                                 padding="same")
         else:
-            self.conv1 = tf.keras.layers.Conv2DTranspose(filters=3,
+            self.conv1 = tf.keras.layers.Conv2DTranspose(filters=3 if labels is None else labels,
                                                 kernel_size=(7, 7),
                                                 strides=2,
                                                 padding="same")
@@ -177,7 +178,10 @@ class ResNetTypeI(tf.keras.Model):
             x = self.bn1(x, training=training)
             # x = tf.nn.relu(x)
             x = self.conv1(x)
-            x = tf.nn.tanh(x)
+            if self.labels is None:
+                x = tf.nn.tanh(x)
+            # else:
+            #     x = tf.nn.softmax(x)
 
         return x
 
@@ -224,8 +228,8 @@ class ResNetTypeII(tf.keras.Model):
         return output
 
 
-def resnet_18(reversed=False):
-    return ResNetTypeI(layer_params=[2, 2, 2, 2], reversed=reversed)
+def resnet_18(reversed=False, labels=None):
+    return ResNetTypeI(layer_params=[2, 2, 2, 2], reversed=reversed, labels=labels)
 
 
 def resnet_34():