Upload 4 files

ckpts/best_model.h5

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a6d0361fa140c1dc3b279bcce8107c28b6e10a4e1bc31f770e5b071a44f5f76d
+size 20800096

models/auto_encoder_gray2color.py

@@ -0,0 +1,92 @@
+
+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()
+    autoencoder.compile(optimizer=Adam(learning_rate=7e-5), loss=tf.keras.losses.MeanSquaredError())

requirements.txt

@@ -0,0 +1,5 @@
+numpy==1.26.4
+tensorflow==2.18.0
+opencv-python==4.11.0.86
+scikit-image==0.25.2
+matplotlib==3.7.2