CAE.py

import os
import glob
import math
import random

import numpy as np
import matplotlib.pyplot as plt

from tensorflow.python import keras
from tensorflow.python.keras import backend as K
from tensorflow.python.keras.models import Model,  Sequential
from tensorflow.python.keras.layers import *
from tensorflow.python.keras.preprocessing.image import load_img, img_to_array, array_to_img, ImageDataGenerator

from tensorflow.python.keras.datasets import mnist


(x_train, _), (x_test, _) = mnist.load_data()
x_train = x_train.reshape(-1, 28, 28, 1)
x_test = x_test.reshape(-1, 28, 28, 1)
x_train = x_train/255.
x_test = x_test/255.


def make_masking_noise_data(data_x, percent=0.1):
    size = data_x.shape
    masking = np.random.binomial(n=1, p=percent, size=size)
    return data_x*masking

x_train_masked = make_masking_noise_data(x_train)
x_test_masked = make_masking_noise_data(x_test)


def make_gaussian_noise_data(data_x, scale=0.8):
    gaussian_data_x = data_x + np.random.normal(loc=0, scale=scale, size=data_x.shape)
    gaussian_data_x = np.clip(gaussian_data_x, 0, 1)
    return gaussian_data_x

x_train_gauss = make_gaussian_noise_data(x_train)
x_test_gauss = make_gaussian_noise_data(x_test)


inputs = Input(shape=(28, 28, 1))
x = Conv2D(filters=16, kernel_size=(3, 3), strides=(1, 1), activation='relu', padding='same')(inputs)
x = MaxPooling2D(pool_size=(2, 2), strides=None, padding='valid')(x)
x = Conv2D(filters=8, kernel_size=(3, 3), strides=(1, 1), activation='relu', padding='same')(x)
x = MaxPooling2D(pool_size=(2, 2), strides=None, padding='valid')(x)

x = Conv2D(filters=8, kernel_size=(3, 3), strides=(1, 1), activation='relu', padding='same')(x)
x = UpSampling2D(size=(2, 2))(x)
x = Conv2D(filters=16, kernel_size=(3, 3), strides=(1, 1), activation='relu', padding='same')(x)
x = UpSampling2D(size=(2, 2))(x)
outputs = Conv2D(filters=1, kernel_size=(3, 3), strides=(1, 1), activation='sigmoid', padding='same')(x)
    
autoencoder = Model(inputs, outputs)

autoencoder.compile(optimizer='adam', loss='binary_crossentropy')
print(autoencoder.summary())

initial_weights = autoencoder.get_weights()
autoencoder.fit(x_train_gauss, x_train, epochs=10, batch_size=20, shuffle=True)
gauss_preds = autoencoder.predict(x_test_gauss)

autoencoder.set_weights(initial_weights)
autoencoder.fit(x_train_masked, x_train, epochs=10, batch_size=20, shuffle=True)
masked_preds = autoencoder.predict(x_test_masked)