write simole cnn autoencoder for watermark removal in python

main.py
# Import libraries
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers

# Define the autoencoder model
def autoencoder(input_shape):
    # Encoder
    inputs = keras.Input(shape=input_shape)
    x = layers.Conv2D(64, 3, activation="relu", padding="same")(inputs)
    x = layers.MaxPooling2D(2, padding="same")(x)
    x = layers.Conv2D(32, 3, activation="relu", padding="same")(x)
    x = layers.MaxPooling2D(2, padding="same")(x)

    # Decoder
    x = layers.Conv2DTranspose(32, 3, strides=2, activation="relu", padding="same")(x)
    x = layers.Conv2DTranspose(64, 3, strides=2, activation="relu", padding="same")(x)
    output = layers.Conv2D(1, 3, activation="sigmoid", padding="same")(x)

    # Create the autoencoder model
    model = keras.Model(inputs, output)
    return model
823 chars
23 lines

main.py
# Load the data
(x_train, _), (x_test, _) = keras.datasets.mnist.load_data()

# Preprocess the data
x_train = x_train.reshape((x_train.shape[0], 28, 28, 1)).astype("float32") / 255.0
x_test = x_test.reshape((x_test.shape[0], 28, 28, 1)).astype("float32") / 255.0

# Add watermark to training data
x_train_watermarked = x_train + 0.5 * tf.random.normal(x_train.shape)

# Define the autoencoder model
model = autoencoder(input_shape=(28, 28, 1))

# Train the model
model.compile(optimizer="adam", loss="mse")
model.fit(x_train_watermarked, x_train, epochs=10, batch_size=128, validation_data=(x_test, x_test))

# Remove watermark from test data using autoencoder
x_test_watermarked = x_test + 0.5 * tf.random.normal(x_test.shape)
x_test_denoised = model.predict(x_test_watermarked)
780 chars
21 lines