04-ExtAE-with-MNIST.ipynb

 
 
 [AE4] - Denoiser and classifier model
 Episode 4 : Construction of a denoiser and classifier model


Objectives :

Building a multiple output model, able to denoise and classify
Understanding a more advanced programming model
The calculation needs being important, it is preferable to use a very simple dataset such as MNIST.
The use of a GPU is often indispensable.
What we're going to do :

Defining a multiple output model using Keras procedural programing model
Build the model
Train it
Follow the learning process
Data Terminology :

clean_train, clean_test for noiseless images 
noisy_train, noisy_test for noisy images
class_train, class_test for the classes to which the images belong 
denoised_test for denoised images at the output of the model
classcat_test for class prediction in model output (is a softmax)
classid_test class prediction (ie: argmax of classcat_test)
 Step 1 - Init python stuff
1.1 - Init
 import numpy as np
from skimage import io
import random

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras.callbacks import ModelCheckpoint, TensorBoard

import os,sys
from importlib import reload
import h5py

from modules.MNIST          import MNIST
from modules.ImagesCallback import ImagesCallback

sys.path.append('..')
import fidle.pwk as pwk

run_dir = './run/AE4'
datasets_dir = pwk.init('AE4', run_dir)
 
 1.2 - Parameters
prepared_dataset : Filename of the prepared dataset (Need 400 Mo, but can be in ./data)
dataset_seed : Random seed for shuffling dataset. 'None' mean using /dev/urandom
scale : % of the dataset to use (1. for 100%)
latent_dim : Dimension of the latent space
train_prop : Percentage for train (the rest being for the test)
batch_size : Batch size
epochs : Nb of epochs for training
 prepared_dataset = './data/mnist-noisy.h5'
dataset_seed     = None

scale            = .1

latent_dim       = 10

train_prop       = .8
batch_size       = 128
epochs           = 30
 
 Override parameters (batch mode) - Just forget this cell
 pwk.override('prepared_dataset', 'dataset_seed', 'scale', 'latent_dim')
pwk.override('train_prop', 'batch_size', 'epochs')
 
 Step 2 - Retrieve dataset
With our MNIST class, in one call, we can reload, rescale, shuffle and split our previously saved dataset :-)
 clean_train,clean_test, noisy_train,noisy_test, class_train,class_test = MNIST.reload_prepared_dataset(scale      = scale, 
                                                                                    train_prop = train_prop,
                                                                                    seed       = dataset_seed,
                                                                                    shuffle    = True,
                                                                                    filename=prepared_dataset )
 
 Step 3 - Build models
 Encoder
 inputs    = keras.Input(shape=(28, 28, 1))
x         = layers.Conv2D(32, 3, activation="relu", strides=2, padding="same")(inputs)
x         = layers.Conv2D(64, 3, activation="relu", strides=2, padding="same")(x)
x         = layers.Flatten()(x)
x         = layers.Dense(16, activation="relu")(x)
z         = layers.Dense(latent_dim)(x)

encoder = keras.Model(inputs, z, name="encoder")
# encoder.summary()
 
 Decoder
 inputs  = keras.Input(shape=(latent_dim,))
x       = layers.Dense(7 * 7 * 64, activation="relu")(inputs)
x       = layers.Reshape((7, 7, 64))(x)
x       = layers.Conv2DTranspose(64, 3, activation="relu", strides=2, padding="same")(x)
x       = layers.Conv2DTranspose(32, 3, activation="relu", strides=2, padding="same")(x)
outputs = layers.Conv2DTranspose(1, 3, activation="sigmoid", padding="same")(x)

decoder = keras.Model(inputs, outputs, name="decoder")
# decoder.summary()
 
 AE
 inputs    = keras.Input(shape=(28, 28, 1))

latents   = encoder(inputs)
outputs   = decoder(latents)

ae = keras.Model(inputs,outputs, name='ae')
 
 CNN
 hidden1     = 100
hidden2     = 100

inputs    = keras.Input(shape=(28, 28, 1))

x         = keras.layers.Conv2D(8, (3,3),  activation='relu')(inputs)
x         = keras.layers.MaxPooling2D((2,2))(x)
x         = keras.layers.Dropout(0.2)(x)

x         = keras.layers.Conv2D(16, (3,3), activation='relu')(x)
x         = keras.layers.MaxPooling2D((2,2))(x)
x         = keras.layers.Dropout(0.2)(x)

x         = keras.layers.Flatten()(x)
x         = keras.layers.Dense(100, activation='relu')(x)
x         = keras.layers.Dropout(0.5)(x)

outputs   = keras.layers.Dense(10, activation='softmax')(x)

cnn       = keras.Model(inputs, outputs, name='cnn')
 
 Final model
 inputs    = keras.Input(shape=(28, 28, 1))

denoised = ae(inputs)
classcat = cnn(inputs)

model = keras.Model(inputs, [denoised, classcat])

model.compile(optimizer='rmsprop', 
              loss={'ae':'binary_crossentropy', 'cnn':'sparse_categorical_crossentropy'},
              loss_weights=[1,1],
              metrics={'cnn':'accuracy'} )
 
 Step 4 - Train
20' on a CPU
1'12 on a GPU (V100, IDRIS)
 # ---- Callback : Images
#
pwk.mkdir( run_dir + '/images')
filename = run_dir + '/images/image-{epoch:03d}-{i:02d}.jpg'
callback_images = ImagesCallback(filename, x=clean_test[:5], encoder=encoder,decoder=decoder)

# ---- Callback : Best model
#
pwk.mkdir( run_dir + '/models')
filename = run_dir + '/models/best_model.h5'
callback_bestmodel = tf.keras.callbacks.ModelCheckpoint(filepath=filename, verbose=0, save_best_only=True)

# ---- Callback tensorboard
#
logdir = run_dir + '/logs'
callback_tensorboard = TensorBoard(log_dir=logdir, histogram_freq=1)

# callbacks_list = [callback_images, callback_bestmodel, callback_tensorboard]
callbacks_list = [callback_images, callback_bestmodel]
 
 pwk.chrono_start()

history = model.fit(noisy_train, [clean_train, class_train],
                 batch_size      = batch_size,
                 epochs          = epochs,
                 verbose         = 1,
                 validation_data = (noisy_test, [clean_test, class_test]),
                 callbacks       = callbacks_list  )

pwk.chrono_show()
 
 Step 5 - History
 pwk.plot_history(history,  plot={'Loss':['loss', 'ae_loss', 'cnn_loss'],
                                 'Validation loss':['val_loss','val_ae_loss', 'val_cnn_loss'], 
                                 'Accuracy':['cnn_accuracy','val_cnn_accuracy']}, save_as='01-history')
 
 Step 6 - Denoising progress
 imgs=[]
for epoch in range(0,epochs,4):
    for i in range(5):
        filename = run_dir + '/images/image-{epoch:03d}-{i:02d}.jpg'.format(epoch=epoch, i=i)
        img      = io.imread(filename)
        imgs.append(img)      

pwk.subtitle('Real images (clean_test) :')
pwk.plot_images(clean_test[:5], None, indices='all', columns=5, x_size=2,y_size=2, interpolation=None, save_as='02-original-real')

pwk.subtitle('Noisy images (noisy_test) :')
pwk.plot_images(noisy_test[:5], None, indices='all', columns=5, x_size=2,y_size=2, interpolation=None, save_as='03-original-noisy')

pwk.subtitle('Evolution during the training period (denoised_test) :')
pwk.plot_images(imgs, None, indices='all', columns=5, x_size=2,y_size=2, interpolation=None, y_padding=0.1, save_as='04-learning')

pwk.subtitle('Noisy images (noisy_test) :')
pwk.plot_images(noisy_test[:5], None, indices='all', columns=5, x_size=2,y_size=2, interpolation=None, save_as=None)

pwk.subtitle('Real images (clean_test) :')
pwk.plot_images(clean_test[:5], None, indices='all', columns=5, x_size=2,y_size=2, interpolation=None, save_as=None)
 
 Step 7 - Evaluation
Note : We will use the following data:\
clean_train, clean_test for noiseless images \
noisy_train, noisy_test for noisy images\
class_train, class_test for the classes to which the images belong \
denoised_test for denoised images at the output of the model\
classcat_test for class prediction in model output (is a softmax)\
classid_test class prediction (ie: argmax of classcat_test)
7.1 - Reload our best model
 model = keras.models.load_model(f'{run_dir}/models/best_model.h5')
 
 7.2 - Let's make a prediction
Note that our model will returns 2 outputs : denoised images from output 1 and class prediction from output 2
 denoised_test, classcat_test = model.predict(noisy_test)

print('Denoised images   (denoised_test) shape : ',denoised_test.shape)
print('Predicted classes (classcat_test) shape : ',classcat_test.shape)
 
 7.3 - Denoised images
 i=random.randint(0,len(denoised_test)-8)
j=i+8

pwk.subtitle('Noisy test images (input):')
pwk.plot_images(noisy_test[i:j], None, indices='all', columns=8, x_size=2,y_size=2, interpolation=None, save_as='05-test-noisy')

pwk.subtitle('Denoised images (output):')
pwk.plot_images(denoised_test[i:j], None, indices='all', columns=8, x_size=2,y_size=2, interpolation=None, save_as='06-test-predict')

pwk.subtitle('Real test images :')
pwk.plot_images(clean_test[i:j], None, indices='all', columns=8, x_size=2,y_size=2, interpolation=None, save_as='07-test-real')
 
 7.4 - Class prediction
Note: The evaluation requires the noisy images as input (noisy_test) and the 2 expected outputs:

the images without noise (clean_test)
the classes (class_test)
 score = model.evaluate(noisy_test, [clean_test, class_test], verbose=0)

pwk.subtitle("Accuracy :")
print(f'Classification accuracy : {score[3]:4.4f}')

pwk.subtitle("Few examples :")
classid_test  = np.argmax(classcat_test, axis=-1)
pwk.plot_images(noisy_test, class_test, range(0,200), columns=12, x_size=1, y_size=1, y_pred=classid_test, save_as='04-predictions')
 
 pwk.end()