Update models/autoencoder_gray2color.py

models/autoencoder_gray2color.py

@@ -1,92 +1,87 @@
-
-import tensorflow as tf
-from tensorflow.keras.layers import Input, Conv2D, MaxPooling2D, UpSampling2D, BatchNormalization, Add, Concatenate, Multiply
-from tensorflow.keras.models import Model
-from tensorflow.keras.optimizers import Adam
-
-# Spatial Attention Layer
-# Define SpatialAttention layer
-class SpatialAttention(tf.keras.layers.Layer):
-    def __init__(self, kernel_size=7, **kwargs):
-        super(SpatialAttention, self).__init__(**kwargs)
-        self.kernel_size = kernel_size
-        self.conv = Conv2D(filters=1, kernel_size=kernel_size, padding='same', activation='sigmoid')
-
-    def call(self, inputs):
-        avg_pool = tf.reduce_mean(inputs, axis=-1, keepdims=True)
-        max_pool = tf.reduce_max(inputs, axis=-1, keepdims=True)
-        concat = Concatenate()([avg_pool, max_pool])
-        attention = self.conv(concat)
-        return Multiply()([inputs, attention])
-
-    def get_config(self):
-        config = super(SpatialAttention, self).get_config()
-        config.update({'kernel_size': self.kernel_size})
-        return config
-        
-# Build Autoencoder
-def build_autoencoder(height, width,):
-    input_img = Input(shape=(height, width, 1))
-    
-    # Encoder
-    x = Conv2D(96, (3, 3), activation='relu', padding='same')(input_img)
-    x = BatchNormalization()(x)
-    x = SpatialAttention()(x)
-    x = MaxPooling2D((2, 2), padding='same')(x)
-    
-    # Residual Block 1
-    residual = Conv2D(192, (1, 1), padding='same')(x)
-    x = Conv2D(192, (3, 3), activation='relu', padding='same')(x)
-    x = BatchNormalization()(x)
-    x = Conv2D(192, (3, 3), activation='relu', padding='same')(x)
-    x = BatchNormalization()(x)
-    x = Add()([x, residual])
-    x = SpatialAttention()(x)
-    x = MaxPooling2D((2, 2), padding='same')(x)
-    
-    # Residual Block 2
-    residual = Conv2D(384, (1, 1), padding='same')(x)
-    x = Conv2D(384, (3, 3), activation='relu', padding='same')(x)
-    x = BatchNormalization()(x)
-    x = Conv2D(384, (3, 3), activation='relu', padding='same')(x)
-    x = BatchNormalization()(x)
-    x = Add()([x, residual])
-    x = SpatialAttention()(x)
-    encoded = MaxPooling2D((2, 2), padding='same')(x)
-    
-    # Decoder
-    x = Conv2D(384, (3, 3), activation='relu', padding='same')(encoded)
-    x = BatchNormalization()(x)
-    x = SpatialAttention()(x)
-    x = UpSampling2D((2, 2))(x)
-    
-    # Residual Block 3
-    residual = Conv2D(192, (1, 1), padding='same')(x)
-    x = Conv2D(192, (3, 3), activation='relu', padding='same')(x)
-    x = BatchNormalization()(x)
-    x = Conv2D(192, (3, 3), activation='relu', padding='same')(x)
-    x = BatchNormalization()(x)
-    x = Add()([x, residual])
-    x = SpatialAttention()(x)
-    x = UpSampling2D((2, 2))(x)
-    
-    x = Conv2D(96, (3, 3), activation='relu', padding='same')(x)
-    x = BatchNormalization()(x)
-    x = SpatialAttention()(x)
-    x = UpSampling2D((2, 2))(x)
-    
-    decoded = Conv2D(2, (3, 3), activation=None, padding='same')(x)
-    
-    return Model(input_img, decoded)
-
-
-
-
-
-if __name__ == "__main__":
-    # Define constants
-    HEIGHT, WIDTH = 512, 512
-    # Compile model
-    autoencoder = build_autoencoder()
-    autoencoder.summary()
+import tensorflow as tf
+from tensorflow.keras.layers import Input, Conv2D, MaxPooling2D, UpSampling2D, BatchNormalization, Add, Concatenate, Multiply
+from tensorflow.keras.models import Model
+from tensorflow.keras.optimizers import Adam
+
+# Set float32 policy
+tf.keras.mixed_precision.set_global_policy('float32')
+
+# Spatial Attention Layer
+class SpatialAttention(tf.keras.layers.Layer):
+    def __init__(self, kernel_size=7, **kwargs):
+        super(SpatialAttention, self).__init__(**kwargs)
+        self.kernel_size = kernel_size
+        self.conv = Conv2D(filters=1, kernel_size=kernel_size, padding='same', activation='sigmoid')
+
+    def call(self, inputs):
+        avg_pool = tf.reduce_mean(inputs, axis=-1, keepdims=True)
+        max_pool = tf.reduce_max(inputs, axis=-1, keepdims=True)
+        concat = Concatenate()([avg_pool, max_pool])
+        attention = self.conv(concat)
+        return Multiply()([inputs, attention])
+
+    def get_config(self):
+        config = super(SpatialAttention, self).get_config()
+        config.update({'kernel_size': self.kernel_size})
+        return config
+
+# Build Autoencoder
+def build_autoencoder(height, width):
+    input_img = Input(shape=(height, width, 1))
+    # Encoder
+    x = Conv2D(96, (3, 3), activation='relu', padding='same')(input_img)
+    x = BatchNormalization()(x)
+    x = SpatialAttention()(x)
+    x = MaxPooling2D((2, 2), padding='same')(x)
+    
+    # Residual Block 1
+    residual = Conv2D(192, (1, 1), padding='same')(x)
+    x = Conv2D(192, (3, 3), activation='relu', padding='same')(x)
+    x = BatchNormalization()(x)
+    x = Conv2D(192, (3, 3), activation='relu', padding='same')(x)
+    x = BatchNormalization()(x)
+    x = Add()([x, residual])
+    x = SpatialAttention()(x)
+    x = MaxPooling2D((2, 2), padding='same')(x)
+    
+    # Residual Block 2
+    residual = Conv2D(384, (1, 1), padding='same')(x)
+    x = Conv2D(384, (3, 3), activation='relu', padding='same')(x)
+    x = BatchNormalization()(x)
+    x = Conv2D(384, (3, 3), activation='relu', padding='same')(x)
+    x = BatchNormalization()(x)
+    x = Add()([x, residual])
+    x = SpatialAttention()(x)
+    encoded = MaxPooling2D((2, 2), padding='same')(x)
+    
+    # Decoder
+    x = Conv2D(384, (3, 3), activation='relu', padding='same')(encoded)
+    x = BatchNormalization()(x)
+    x = SpatialAttention()(x)
+    x = UpSampling2D((2, 2))(x)
+    
+    # Residual Block 3
+    residual = Conv2D(192, (1, 1), padding='same')(x)
+    x = Conv2D(192, (3, 3), activation='relu', padding='same')(x)
+    x = BatchNormalization()(x)
+    x = Conv2D(192, (3, 3), activation='relu', padding='same')(x)
+    x = BatchNormalization()(x)
+    x = Add()([x, residual])
+    x = SpatialAttention()(x)
+    x = UpSampling2D((2, 2))(x)
+    
+    x = Conv2D(96, (3, 3), activation='relu', padding='same')(x)
+    x = BatchNormalization()(x)
+    x = SpatialAttention()(x)
+    x = UpSampling2D((2, 2))(x)
+    decoded = Conv2D(2, (3, 3), activation=None, padding='same')(x)
+    
+    return Model(input_img, decoded)
+
+if __name__ == "__main__":
+    # Define constants
+    HEIGHT, WIDTH = 512, 512
+    # Compile model
+    autoencoder = build_autoencoder(HEIGHT, WIDTH)
+    autoencoder.summary()
     autoencoder.compile(optimizer=Adam(learning_rate=7e-5), loss=tf.keras.losses.MeanSquaredError())