from matplotlib import pyplot
from numpy import float32, random, array, prod, arange, clip, ndarray
from tensorflow.keras import layers, Model, backend
from tensorflow.keras.datasets import cifar10
from tensorflow.keras.layers import Input, Conv2D, MaxPooling2D, UpSampling2D

data = cifar10.load_data()
(x_train, _), (x_test, _) = data
x_train = array(x_train)
x_test = array(x_test)

x_train.shape

i = 1
pyplot.figure(figsize=(10, 10))
for image in x_train:
    if i > 5:
        break
    pyplot.subplot(1, 5, i)
    pyplot.imshow(image)
    pyplot.axis("off")
    i += 1
pyplot.show()

x_train = x_train.astype(float32) / 255.0
x_test = x_test.astype(float32) / 255.0

i = 1
pyplot.figure(figsize=(10, 10))
for image in x_train:
    if i > 5:
        break
    pyplot.subplot(1, 5, i)
    pyplot.imshow(image)
    pyplot.axis("off")
    i += 1
pyplot.show()

x_train_flattened = x_train.reshape((len(x_train), prod(x_train.shape[1:])))
x_test_flattened = x_test.reshape((len(x_test), prod(x_test.shape[1:])))

x_train_flattened.shape, x_test_flattened.shape

encoding_dim = 100  # This is the size of our encoded representations

Create Autoencoder Model

input_img = Input(shape=(3072,))

# Encoder
x = layers.Dense(1000, activation="relu")(input_img)
x = layers.Dense(500, activation="relu")(x)
x = layers.Dense(200, activation="relu")(x)
encoded = layers.Dense(encoding_dim, activation="relu")(x)  # Define encoding dimension

# Decoder
x = layers.Dense(200, activation="relu")(encoded)
x = layers.Dense(500, activation="relu")(x)
x = layers.Dense(1000, activation="relu")(x)
decoded = layers.Dense(3072, activation="sigmoid")(x)  # Match output dimension to input

autoencoder = Model(input_img, decoded)

autoencoder.compile(loss="mse", optimizer="adam")

h = autoencoder.fit(x_train_flattened, x_train_flattened, epochs=100, batch_size=64)

pyplot.figure(figsize=(10, 10))
pyplot.plot(arange(0, 100), h.history["loss"])
pyplot.title("Training Loss")
pyplot.xlabel("Epoch")
pyplot.ylabel("Loss")
pyplot.legend()
pyplot.show()

encoder = Model(input_img, encoded)

encoded_input = Input(shape=(encoding_dim,))
decoder_l1 = autoencoder.layers[-4]
decoder_l2 = autoencoder.layers[-3]
decoder_l3 = autoencoder.layers[-2]
decoder_l4 = autoencoder.layers[-1]
# as if I am multiplying the last layer with the code layer to get the last output
decoder = Model(
    encoded_input, decoder_l4(decoder_l3(decoder_l2(decoder_l1(encoded_input))))
)

encoded_images = encoder.predict(x_test_flattened)
decoded_images = decoder.predict(encoded_images)

n = 15
pyplot.figure(figsize=(20, 4))
for i in range(n):
    ax = pyplot.subplot(2, n, i + 1)
    pyplot.imshow(x_test_flattened[i].reshape(32, 32, 3))
    pyplot.gray()
    ax.get_xaxis().set_visible(False)
    ax.get_yaxis().set_visible(False)

    ax = pyplot.subplot(2, n, i + 1 + n)
    pyplot.imshow(decoded_images[i].reshape(32, 32, 3))
    pyplot.gray()
    ax.get_xaxis().set_visible(False)
    ax.get_yaxis().set_visible(False)
pyplot.show()

n = 15
pyplot.figure(figsize=(20, 4))
for i in range(n):
    ax = pyplot.subplot(3, n, i + 1)
    pyplot.imshow(x_test_flattened[i].reshape(32, 32, 3))
    pyplot.gray()
    ax.get_xaxis().set_visible(False)
    ax.get_yaxis().set_visible(False)

    ax = pyplot.subplot(3, n, i + 1 + n)
    pyplot.imshow(encoded_images[i].reshape(10, 10))
    pyplot.gray()
    ax.get_xaxis().set_visible(False)
    ax.get_yaxis().set_visible(False)

    ax = pyplot.subplot(3, n, i + 1 + 2 * n)
    pyplot.imshow(decoded_images[i].reshape(32, 32, 3))
    pyplot.gray()
    ax.get_xaxis().set_visible(False)
    ax.get_yaxis().set_visible(False)
pyplot.show()

def add_gaussian_noise(images: ndarray, mean: float, std_dev: float) -> ndarray:
    noise = random.normal(mean, std_dev, size=images.shape)
    return clip(images + noise, 0.0, 1.0)

mean = 0
sigma = 0.1

x_train_noisy = add_gaussian_noise(x_train, mean, sigma)
x_test_noisy = add_gaussian_noise(x_test, mean, sigma)

x_train_noisy.shape, x_test_noisy.shape

x_train_noisy_flattened = x_train_noisy.reshape(
    (len(x_train_noisy), prod(x_train_noisy.shape[1:]))
)
x_test_noisy_flattened = x_test_noisy.reshape(
    (len(x_test_noisy), prod(x_test_noisy.shape[1:]))
)

backend.clear_session()

After Adding Noise

autoencoder.compile(optimizer="adam", loss="mse")

h = autoencoder.fit(
    x_train_noisy_flattened, x_train_flattened, epochs=100, batch_size=64
)

encoded_images_noisy = encoder.predict(x_test_noisy_flattened)
decoded_images_noisy = decoder.predict(encoded_images_noisy)

n = 15
pyplot.figure(figsize=(20, 4))
for i in range(n):
    ax = pyplot.subplot(3, n, i + 1)
    pyplot.imshow(x_test_noisy[i].reshape(32, 32, 3))
    pyplot.gray()
    ax.get_xaxis().set_visible(False)
    ax.get_yaxis().set_visible(False)

    ax = pyplot.subplot(3, n, i + 1 + n)
    pyplot.imshow(encoded_images_noisy[i].reshape(10, 10))
    pyplot.gray()
    ax.get_xaxis().set_visible(False)
    ax.get_yaxis().set_visible(False)

    ax = pyplot.subplot(3, n, i + 1 + 2 * n)
    pyplot.imshow(decoded_images_noisy[i].reshape(32, 32, 3))
    pyplot.gray()
    ax.get_xaxis().set_visible(False)
    ax.get_yaxis().set_visible(False)
pyplot.show()

def add_noise_to_code(encoded: ndarray, mean: float, std_dev: float) -> ndarray:
    noise = backend.random_normal(
        shape=backend.shape(encoded), mean=mean, stddev=std_dev
    )
    return encoded + noise

encoded_images = encoder.predict(x_test_flattened)
encoded_images_noisy = add_noise_to_code(encoded_images_noisy, mean, sigma)
decoded_images_noisy = decoder.predict(encoded_images_noisy)

encoded_images_noisy = array(encoded_images_noisy)

n = 15
pyplot.figure(figsize=(20, 4))
for i in range(n):
    ax = pyplot.subplot(3, n, i + 1)
    pyplot.imshow(x_test_flattened[i].reshape(32, 32, 3))
    pyplot.gray()
    ax.get_xaxis().set_visible(False)
    ax.get_yaxis().set_visible(False)

    ax = pyplot.subplot(3, n, i + 1 + n)
    pyplot.imshow(encoded_images_noisy[i].reshape(10, 10))
    pyplot.gray()
    ax.get_xaxis().set_visible(False)
    ax.get_yaxis().set_visible(False)

    ax = pyplot.subplot(3, n, i + 1 + 2 * n)
    pyplot.imshow(decoded_images_noisy[i].reshape(32, 32, 3))
    pyplot.gray()
    ax.get_xaxis().set_visible(False)
    ax.get_yaxis().set_visible(False)
pyplot.show()

Create CNN Autoencoder Model

input_img = Input(shape=(32, 32, 3))

# Encoder
x = Conv2D(32, (3, 3), activation="relu", padding="same")(input_img)
x = MaxPooling2D((2, 2), padding="same")(x)
x = Conv2D(64, (3, 3), activation="relu", padding="same")(x)
encoded = MaxPooling2D((2, 2), padding="same")(x)

# Decoder
x = Conv2D(64, (3, 3), activation="relu", padding="same")(encoded)
x = UpSampling2D((2, 2))(x)
x = Conv2D(32, (3, 3), activation="relu", padding="same")(x)
x = UpSampling2D((2, 2))(x)
decoded = Conv2D(3, (3, 3), activation="sigmoid", padding="same")(x)

autoencoder = Model(input_img, decoded)

autoencoder.compile(optimizer="adam", loss="binary_crossentropy")

h = autoencoder.fit(x_train, x_train, epochs=100, batch_size=64)

pyplot.figure(figsize=(10, 10))
pyplot.plot(arange(0, 100), h.history["loss"])
pyplot.title("Training Loss")
pyplot.xlabel("Epoch")
pyplot.ylabel("Loss")
pyplot.legend()
pyplot.show()