Update models/unet_gray2color.py

models/unet_gray2color.py

@@ -14,10 +14,9 @@ from skimage.color import rgb2lab, lab2rgb
 from skimage.metrics import peak_signal_noise_ratio
 import matplotlib.pyplot as plt
 
-# Disable mixed precision to avoid dtype mismatches
+# Disable mixed precision
 tf.keras.mixed_precision.set_global_policy('float32')
 
-# Custom self-attention layer with serialization support
 @tf.keras.utils.register_keras_serializable()
 class SelfAttentionLayer(Layer):
     def __init__(self, num_heads, key_dim, **kwargs):
@@ -28,12 +27,9 @@ class SelfAttentionLayer(Layer):
         self.ln = LayerNormalization(epsilon=1e-6)
 
     def build(self, input_shape):
-        # input_shape: (batch_size, height, width, channels)
         batch_size, height, width, channels = input_shape
-        attention_shape = (batch_size, height * width, channels)  # Shape after reshape
-        # Build MultiHeadAttention with query, key, and value shapes
+        attention_shape = (batch_size, height * width, channels)
         self.mha.build(query_shape=attention_shape, value_shape=attention_shape)
-        # Build LayerNormalization with the original input shape
         self.ln.build(input_shape)
         super(SelfAttentionLayer, self).build(input_shape)
 
@@ -41,7 +37,6 @@ class SelfAttentionLayer(Layer):
         b, h, w, c = tf.shape(x)[0], x.shape[1], x.shape[2], x.shape[3]
         attention_input = tf.reshape(x, [b, h * w, c])
         attention_output = self.mha(attention_input, attention_input)
-        # Cast attention_output to match x's dtype
         attention_output = tf.cast(attention_output, dtype=x.dtype)
         attention_output = tf.reshape(attention_output, [b, h, w, c])
         return self.ln(x + attention_output)
@@ -57,9 +52,8 @@ class SelfAttentionLayer(Layer):
         })
         return config
 
-def attention_unet_model(input_shape=(256, 256, 1)):
+def attention_unet_model(input_shape=(1024, 1024, 1)):
     inputs = Input(input_shape)
-    # Encoder with reduced filters
     c1 = Conv2D(16, (3, 3), activation='relu', padding='same')(inputs)
     c1 = BatchNormalization()(c1)
     c1 = Conv2D(16, (3, 3), activation='relu', padding='same')(c1)
@@ -78,14 +72,12 @@ def attention_unet_model(input_shape=(256, 256, 1)):
     c3 = BatchNormalization()(c3)
     p3 = MaxPooling2D((2, 2))(c3)
     
-    # Bottleneck with reduced filters and attention
     c4 = Conv2D(128, (3, 3), activation='relu', padding='same')(p3)
     c4 = BatchNormalization()(c4)
     c4 = Conv2D(128, (3, 3), activation='relu', padding='same')(c4)
     c4 = BatchNormalization()(c4)
     c4 = SelfAttentionLayer(num_heads=2, key_dim=32)(c4)
     
-    # Attention gate
     def attention_gate(g, s, num_filters):
         g_conv = Conv2D(num_filters, (1, 1), padding='same')(g)
         s_conv = Conv2D(num_filters, (1, 1), padding='same')(s)
@@ -94,7 +86,6 @@ def attention_unet_model(input_shape=(256, 256, 1)):
         attn = Conv2D(1, (1, 1), padding='same', activation='sigmoid')(attn)
         return s * attn
     
-    # Decoder with reduced filters
     u5 = UpSampling2D((2, 2))(c4)
     a5 = attention_gate(u5, c3, 64)
     u5 = Concatenate()([u5, a5])
@@ -119,15 +110,12 @@ def attention_unet_model(input_shape=(256, 256, 1)):
     c7 = Conv2D(16, (3, 3), activation='relu', padding='same')(c7)
     c7 = BatchNormalization()(c7)
     
-    # Output layer
     outputs = Conv2D(2, (1, 1), activation='tanh', padding='same')(c7)
     model = Model(inputs, outputs)
     return model
 
 if __name__ == "__main__":
-    # Define constants
-    HEIGHT, WIDTH = 256, 256  # Match function definition
-    # Compile model
+    HEIGHT, WIDTH = 1024, 1024
     model = attention_unet_model(input_shape=(HEIGHT, WIDTH, 1))
     model.summary()
     model.compile(optimizer=Adam(learning_rate=7e-5), loss=tf.keras.losses.MeanSquaredError())