Update models/transformer_gray2color.py

models/transformer_gray2color.py

@@ -3,20 +3,16 @@ from tensorflow.keras.layers import Input, Dense, LayerNormalization, Dropout, M
 from tensorflow.keras.models import Model
 from tensorflow.keras.optimizers import Adam
 
-# Set float32 policy
 tf.keras.mixed_precision.set_global_policy('float32')
 
-# Define Transformer model
 def transformer_model(input_shape=(1024, 1024, 1), patch_size=8, d_model=32, num_heads=4, ff_dim=64, num_layers=2, dropout_rate=0.1):
     HEIGHT, WIDTH, _ = input_shape
     num_patches = (HEIGHT // patch_size) * (WIDTH // patch_size)
     
     inputs = Input(shape=input_shape)
-    # Patch extraction
     x = Conv2D(d_model, (patch_size, patch_size), strides=(patch_size, patch_size), padding='valid')(inputs)
     x = Reshape((num_patches, d_model))(x)
     
-    # Transformer layers
     for _ in range(num_layers):
         attn_output = MultiHeadAttention(num_heads=num_heads, key_dim=d_model // num_heads)(x, x)
         attn_output = Dropout(dropout_rate)(attn_output)
@@ -28,7 +24,6 @@ def transformer_model(input_shape=(1024, 1024, 1), patch_size=8, d_model=32, num
         x = Add()([x, ff_output])
         x = LayerNormalization(epsilon=1e-6)(x)
     
-    # Decoder: Reconstruct image
     x = Dense(2)(x)
     x = Reshape((HEIGHT // patch_size, WIDTH // patch_size, 2))(x)
     x = UpSampling2D(size=(patch_size, patch_size), interpolation='bilinear')(x)
@@ -37,9 +32,7 @@ def transformer_model(input_shape=(1024, 1024, 1), patch_size=8, d_model=32, num
     return Model(inputs, outputs)
 
 if __name__ == "__main__":
-    # Define constants
     HEIGHT, WIDTH = 1024, 1024
-    # Instantiate and compile the model
     model = transformer_model(input_shape=(HEIGHT, WIDTH, 1), patch_size=8, d_model=32, num_heads=4, ff_dim=64, num_layers=2)
     model.summary()
     model.compile(optimizer=Adam(learning_rate=7e-5), loss=tf.keras.losses.MeanSquaredError())