AE/05-ExtAE-with-MNIST.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img width=\"800px\" src=\"../fidle/img/00-Fidle-header-01.svg\"></img>\n",
    "\n",
    "# <!-- TITLE --> [AE5] - Advanced denoiser and classifier model\n",
    "<!-- DESC --> Episode 5 : Construction of an advanced denoiser and classifier model\n",
    "\n",
    "<!-- AUTHOR : Jean-Luc Parouty (CNRS/SIMaP) -->\n",
    "\n",
    "## Objectives :\n",
    " - Building a multiple output model, able to **denoise** and **classify**\n",
    " - Understanding a more complex **advanced programming model**\n",
    "\n",
    "The calculation needs being important, it is preferable to use a very simple dataset such as MNIST.  \n",
    "The use of a GPU is often indispensable.\n",
    "\n",
    "## What we're going to do :\n",
    "\n",
    " - Defining a multiple output model using Keras procedural programing model\n",
    " - Build the model\n",
    " - Train it\n",
    " - Follow the learning process\n",
    " \n",
    "## Data Terminology :\n",
    "- `clean_train`, `clean_test` for noiseless images \n",
    "- `noisy_train`, `noisy_test` for noisy images\n",
    "- `class_train`, `class_test` for the classes to which the images belong \n",
    "- `denoised_test` for denoised images at the output of the model\n",
    "- `classcat_test` for class prediction in model output (is a softmax)\n",
    "- `classid_test` class prediction (ie: argmax of classcat_test)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 1 - Init python stuff\n",
    "### 1.1 - Init"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from skimage import io\n",
    "import random\n",
    "\n",
    "import tensorflow as tf\n",
    "from tensorflow import keras\n",
    "from tensorflow.keras import layers\n",
    "from tensorflow.keras.callbacks import ModelCheckpoint, TensorBoard\n",
    "\n",
    "import os,sys\n",
    "from importlib import reload\n",
    "import h5py\n",
    "\n",
    "from modules.MNIST          import MNIST\n",
    "from modules.ImagesCallback import ImagesCallback\n",
    "\n",
    "sys.path.append('..')\n",
    "import fidle.pwk as pwk\n",
    "\n",
    "run_dir = './run/AE5'\n",
    "datasets_dir = pwk.init('AE5', run_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.2 - Parameters\n",
    "`prepared_dataset` : Filename of the prepared dataset (Need 400 Mo, but can be in ./data)  \n",
    "`dataset_seed` : Random seed for shuffling dataset  \n",
    "`scale` : % of the dataset to use (1. for 100%)  \n",
    "`latent_dim` : Dimension of the latent space  \n",
    "`train_prop` : Percentage for train (the rest being for the test)\n",
    "`batch_size` : Batch size  \n",
    "`epochs` : Nb of epochs for training\\\n",
    "`fit_verbosity` is the verbosity during training : 0 = silent, 1 = progress bar, 2 = one line per epoch\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prepared_dataset = './data/mnist-noisy.h5'\n",
    "dataset_seed     = None\n",
    "\n",
    "scale            = 1\n",
    "\n",
    "latent_dim       = 10\n",
    "\n",
    "train_prop       = .8\n",
    "batch_size       = 128\n",
    "epochs           = 20\n",
    "fit_verbosity    = 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Override parameters (batch mode) - Just forget this cell"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pwk.override('prepared_dataset', 'dataset_seed', 'scale', 'latent_dim')\n",
    "pwk.override('train_prop', 'batch_size', 'epochs')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 2 - Retrieve dataset\n",
    "With our MNIST class, in one call, we can reload, rescale, shuffle and split our previously saved dataset :-)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "clean_train,clean_test, noisy_train,noisy_test, class_train,class_test = MNIST.reload_prepared_dataset(scale      = scale, \n",
    "                                                                                    train_prop = train_prop,\n",
    "                                                                                    seed       = dataset_seed,\n",
    "                                                                                    shuffle    = True,\n",
    "                                                                                    filename=prepared_dataset )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 3 - Build model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Encoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs    = keras.Input(shape=(28, 28, 1))\n",
    "x         = layers.Conv2D(32, 3, activation=\"relu\", strides=2, padding=\"same\")(inputs)\n",
    "x         = layers.Conv2D(64, 3, activation=\"relu\", strides=2, padding=\"same\")(x)\n",
    "x         = layers.Flatten()(x)\n",
    "x         = layers.Dense(16, activation=\"relu\")(x)\n",
    "z         = layers.Dense(latent_dim)(x)\n",
    "\n",
    "encoder = keras.Model(inputs, z, name=\"encoder\")\n",
    "# encoder.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Decoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs  = keras.Input(shape=(latent_dim,))\n",
    "x       = layers.Dense(7 * 7 * 64, activation=\"relu\")(inputs)\n",
    "x       = layers.Reshape((7, 7, 64))(x)\n",
    "x       = layers.Conv2DTranspose(64, 3, activation=\"relu\", strides=2, padding=\"same\")(x)\n",
    "x       = layers.Conv2DTranspose(32, 3, activation=\"relu\", strides=2, padding=\"same\")(x)\n",
    "outputs = layers.Conv2DTranspose(1, 3, activation=\"sigmoid\", padding=\"same\")(x)\n",
    "\n",
    "decoder = keras.Model(inputs, outputs, name=\"decoder\")\n",
    "# decoder.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### AE\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs    = keras.Input(shape=(28, 28, 1))\n",
    "\n",
    "latents   = encoder(inputs)\n",
    "outputs   = decoder(latents)\n",
    "\n",
    "ae = keras.Model(inputs,outputs, name='ae')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### CNN1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs    = keras.Input(shape=(28, 28, 1))\n",
    "\n",
    "x         = keras.layers.Conv2D(8, (3,3),  activation='relu')(inputs)\n",
    "x         = keras.layers.MaxPooling2D((2,2))(x)\n",
    "x         = keras.layers.Dropout(0.2)(x)\n",
    "\n",
    "x         = keras.layers.Conv2D(16, (3,3), activation='relu')(x)\n",
    "x         = keras.layers.MaxPooling2D((2,2))(x)\n",
    "x         = keras.layers.Dropout(0.2)(x)\n",
    "\n",
    "x         = keras.layers.Flatten()(x)\n",
    "x         = keras.layers.Dense(100, activation='relu')(x)\n",
    "outputs   = keras.layers.Dropout(0.5)(x)\n",
    "\n",
    "cnn1       = keras.Model(inputs, outputs, name='cnn1')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### CNN2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs    = keras.Input(shape=(28, 28, 1))\n",
    "\n",
    "x         = keras.layers.Conv2D(32, (5,5),  activation='relu')(inputs)\n",
    "x         = keras.layers.MaxPooling2D((2,2))(x)\n",
    "x         = keras.layers.Dropout(0.3)(x)\n",
    "\n",
    "x         = keras.layers.Conv2D(64, (5,5), activation='relu')(x)\n",
    "x         = keras.layers.MaxPooling2D((2,2))(x)\n",
    "x         = keras.layers.Dropout(0.3)(x)\n",
    "\n",
    "x         = keras.layers.Flatten()(x)\n",
    "x         = keras.layers.Dense(50, activation='relu')(x)\n",
    "outputs   = keras.layers.Dropout(0.3)(x)\n",
    "\n",
    "cnn2       = keras.Model(inputs, outputs, name='cnn2')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Final model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs    = keras.Input(shape=(28, 28, 1))\n",
    "\n",
    "denoised = ae(inputs)\n",
    "\n",
    "branch_1 = cnn1(inputs)\n",
    "branch_2 = cnn2(inputs)\n",
    "\n",
    "x        = keras.layers.concatenate([branch_1,branch_2], axis=1)\n",
    "\n",
    "classcat = keras.layers.Dense(10, activation='softmax', name='cnn')(x)\n",
    "\n",
    "\n",
    "model = keras.Model(inputs, [denoised, classcat])\n",
    "\n",
    "model.compile(optimizer='rmsprop', \n",
    "              loss={'ae':'binary_crossentropy', 'cnn':'sparse_categorical_crossentropy'},\n",
    "              loss_weights=[1,1],\n",
    "              metrics={'cnn':'accuracy'} )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 4 - Train\n",
    "20' on a CPU  \n",
    "1'30 on a GPU (V100, IDRIS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ---- Callback : Images\n",
    "#\n",
    "pwk.mkdir( run_dir + '/images')\n",
    "filename = run_dir + '/images/image-{epoch:03d}-{i:02d}.jpg'\n",
    "callback_images = ImagesCallback(filename, x=clean_test[:5], encoder=encoder,decoder=decoder)\n",
    "\n",
    "# ---- Callback : Best model\n",
    "#\n",
    "pwk.mkdir( run_dir + '/models')\n",
    "filename = run_dir + '/models/best_model.h5'\n",
    "callback_bestmodel = tf.keras.callbacks.ModelCheckpoint(filepath=filename, verbose=0, save_best_only=True)\n",
    "\n",
    "# ---- Callback tensorboard\n",
    "#\n",
    "logdir = run_dir + '/logs'\n",
    "callback_tensorboard = TensorBoard(log_dir=logdir, histogram_freq=1)\n",
    "\n",
    "# callbacks_list = [callback_images, callback_bestmodel, callback_tensorboard]\n",
    "callbacks_list = [callback_images, callback_bestmodel]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pwk.chrono_start()\n",
    "\n",
    "history = model.fit(noisy_train, [clean_train, class_train],\n",
    "                 batch_size      = batch_size,\n",
    "                 epochs          = epochs,\n",
    "                 verbose         = fit_verbosity,\n",
    "                 validation_data = (noisy_test, [clean_test, class_test]),\n",
    "                 callbacks       = callbacks_list  )\n",
    "\n",
    "pwk.chrono_show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 5 - History"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pwk.plot_history(history,  plot={'Loss':['loss', 'ae_loss', 'cnn_loss'],\n",
    "                                 'Validation loss':['val_loss','val_ae_loss', 'val_cnn_loss'], \n",
    "                                 'Accuracy':['cnn_accuracy','val_cnn_accuracy']}, save_as='01-history')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 6 - Denoising progress"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "imgs=[]\n",
    "for epoch in range(0,epochs,4):\n",
    "    for i in range(5):\n",
    "        filename = run_dir + '/images/image-{epoch:03d}-{i:02d}.jpg'.format(epoch=epoch, i=i)\n",
    "        img      = io.imread(filename)\n",
    "        imgs.append(img)      \n",
    "\n",
    "pwk.subtitle('Real images (clean_test) :')\n",
    "pwk.plot_images(clean_test[:5], None, indices='all', columns=5, x_size=2,y_size=2, interpolation=None, save_as='02-original-real')\n",
    "\n",
    "pwk.subtitle('Noisy images (noisy_test) :')\n",
    "pwk.plot_images(noisy_test[:5], None, indices='all', columns=5, x_size=2,y_size=2, interpolation=None, save_as='03-original-noisy')\n",
    "\n",
    "pwk.subtitle('Evolution during the training period (denoised_test) :')\n",
    "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')\n",
    "\n",
    "pwk.subtitle('Noisy images (noisy_test) :')\n",
    "pwk.plot_images(noisy_test[:5], None, indices='all', columns=5, x_size=2,y_size=2, interpolation=None, save_as=None)\n",
    "\n",
    "pwk.subtitle('Real images (clean_test) :')\n",
    "pwk.plot_images(clean_test[:5], None, indices='all', columns=5, x_size=2,y_size=2, interpolation=None, save_as=None)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 7 - Evaluation\n",
    "**Note :** We will use the following data:\\\n",
    "`clean_train`, `clean_test` for noiseless images \\\n",
    "`noisy_train`, `noisy_test` for noisy images\\\n",
    "`class_train`, `class_test` for the classes to which the images belong \\\n",
    "`denoised_test` for denoised images at the output of the model\\\n",
    "`classcat_test` for class prediction in model output (is a softmax)\\\n",
    "`classid_test` class prediction (ie: argmax of classcat_test)\n",
    " \n",
    "### 7.1 - Reload our best model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = keras.models.load_model(f'{run_dir}/models/best_model.h5')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 7.2 - Let's make a prediction\n",
    "Note that our model will returns 2 outputs : **denoised images** from output 1 and **class prediction** from output 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "denoised_test, classcat_test = model.predict(noisy_test)\n",
    "\n",
    "print('Denoised images   (denoised_test) shape : ',denoised_test.shape)\n",
    "print('Predicted classes (classcat_test) shape : ',classcat_test.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 7.3 - Denoised images "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "i=random.randint(0,len(denoised_test)-8)\n",
    "j=i+8\n",
    "\n",
    "pwk.subtitle('Noisy test images (input):')\n",
    "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')\n",
    "\n",
    "pwk.subtitle('Denoised images (output):')\n",
    "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')\n",
    "\n",
    "pwk.subtitle('Real test images :')\n",
    "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')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 7.4 - Class prediction\n",
    "Note: The evaluation requires the noisy images as input (noisy_test) and the 2 expected outputs:\n",
    " - the images without noise (clean_test)\n",
    " - the classes (class_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "score = model.evaluate(noisy_test, [clean_test, class_test], verbose=0)\n",
    "\n",
    "pwk.subtitle(\"Accuracy :\")\n",
    "print(f'Classification accuracy : {score[3]:4.4f}')\n",
    "\n",
    "pwk.subtitle(\"Few examples :\")\n",
    "classid_test  = np.argmax(classcat_test, axis=-1)\n",
    "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')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pwk.end()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "<img width=\"80px\" src=\"../fidle/img/00-Fidle-logo-01.svg\"></img>"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}