train_RAMS_network.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train RAMS Deep Neural Network on Proba-V Dataset\n",
    "![proba_v_dataset](media/RAMS.png \"Logo Title Text 1\")\n",
    "\n",
    "The following notebook provides a script to train the residual attention network for multi-image super-resolution (RAMS). It makes use of the pre-processed dataset (train and validation) saved in the 'dataset' folder and using the main settings it selects a band to train with. \n",
    "\n",
    "**NB**: We strongly discouraged to run this notebook without an available GPU on the host machine. The original training (ckpt folder) has been performed on a 2080 Ti GPU card with 11GB of memory in approximately 24 hours.\n",
    "\n",
    "**The notebook is divided in**:\n",
    "- 1.0 [Dataset Loading](#loading)\n",
    "- 2.0 [Dataset Pre-Processing](#preprocessing)\n",
    "    - 2.1 Make patches\n",
    "    - 2.2 Clarity patches check\n",
    "    - 2.3 Pre-augment dataset (temporal permutation)\n",
    "- 3.0 [Build the network](#network)\n",
    "- 4.0 [Train the network](#train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-03-06T14:36:07.882578Z",
     "start_time": "2020-03-06T14:36:07.757190Z"
    }
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "%reload_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-03-06T14:36:09.065498Z",
     "start_time": "2020-03-06T14:36:08.039481Z"
    }
   },
   "outputs": [],
   "source": [
    "# import utils and basic libraries\n",
    "import os\n",
    "from tqdm.notebook import tqdm\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "from utils.preprocessing import gen_sub, bicubic\n",
    "from utils.loss import l1_loss, psnr, ssim\n",
    "from utils.network import RAMS\n",
    "from utils.training import Trainer\n",
    "from skimage import io\n",
    "from zipfile import ZipFile"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-03-06T14:36:09.183877Z",
     "start_time": "2020-03-06T14:36:09.149475Z"
    }
   },
   "outputs": [],
   "source": [
    "# gpu settings (we strongly discouraged to run this notebook without an available GPU)\n",
    "gpus = tf.config.experimental.list_physical_devices('GPU')\n",
    "tf.config.experimental.set_visible_devices(gpus[0], 'GPU')\n",
    "tf.config.experimental.set_memory_growth(gpus[0], True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-03-06T14:36:09.689561Z",
     "start_time": "2020-03-06T14:36:09.671857Z"
    }
   },
   "outputs": [],
   "source": [
    "#-------------\n",
    "# General Settings\n",
    "#-------------\n",
    "PATH_DATASET = 'dataset' # pre-processed dataset path\n",
    "name_net = 'RAMS' # name of the network\n",
    "LR_SIZE = 32 # pathces dimension\n",
    "SCALE = 3 # upscale of the proba-v dataset is 3\n",
    "HR_SIZE = LR_SIZE * SCALE # upscale of the dataset is 3\n",
    "OVERLAP = 32 # overlap between pathces\n",
    "CLEAN_PATH_PX = 0.85 # percentage of clean pixels to accept a patch\n",
    "band = 'NIR' # choose the band for the training\n",
    "checkpoint_dir = f'ckpt/{band}_{name_net}_retrain' # weights path\n",
    "log_dir = 'logs' # tensorboard logs path\n",
    "submission_dir = 'submission' # submission dir"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#-------------\n",
    "# Network Settings\n",
    "#-------------\n",
    "FILTERS = 32 # features map in the network\n",
    "KERNEL_SIZE = 3 # convolutional kernel size dimension (either 3D and 2D)\n",
    "CHANNELS = 9 # number of temporal steps\n",
    "R = 8 # attention compression\n",
    "N = 12 # number of residual feature attention blocks\n",
    "lr = 1e-4 # learning rate (Nadam optimizer)\n",
    "BATCH_SIZE = 32 # batch size\n",
    "EPOCHS_N = 100 # number of epochs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create logs folder\n",
    "if not os.path.exists(log_dir):\n",
    "    os.mkdir(log_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"loading\"></a>\n",
    "# 1.0 Dataset loading"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load training dataset\n",
    "X_train = np.load(os.path.join(PATH_DATASET, f'X_{band}_train.npy'))\n",
    "y_train = np.load(os.path.join(PATH_DATASET, f'y_{band}_train.npy'))\n",
    "y_train_mask = np.load(os.path.join(PATH_DATASET, f'y_{band}_train_masks.npy'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load validation dataset\n",
    "X_val = np.load(os.path.join(PATH_DATASET, f'X_{band}_val.npy'))\n",
    "y_val = np.load(os.path.join(PATH_DATASET, f'y_{band}_val.npy'))\n",
    "y_val_mask = np.load(os.path.join(PATH_DATASET, f'y_{band}_val_masks.npy'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# print loaded dataset info\n",
    "print('X_train: ', X_train.shape)\n",
    "print('y_train: ', y_train.shape)\n",
    "print('y_train_mask: ', y_train_mask.shape)\n",
    "\n",
    "\n",
    "print('X_val: ', X_val.shape)\n",
    "print('y_val: ', y_val.shape)\n",
    "print('y_val_mask: ', y_val_mask.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"preprocessing\"></a>\n",
    "# 2.0 Dataset Pre-Processing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.1 Make patches"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create patches for LR images\n",
    "d = LR_SIZE  # 32x32 patches\n",
    "s = OVERLAP  # overlapping patches\n",
    "# Ex: n = (128-d)/s+1 = 7 -> 49 sub images from each image\n",
    "\n",
    "X_train_patches = gen_sub(X_train,d,s)\n",
    "X_val_patches = gen_sub(X_val,d,s)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create patches for HR images and masks\n",
    "d = HR_SIZE  # 96x96 patches\n",
    "s = OVERLAP * SCALE  # overlapping patches\n",
    "# Ex: n = (384-d)/s+1 = 7 -> 49 sub images from each image\n",
    "\n",
    "y_train_patches = gen_sub(y_train,d,s)\n",
    "y_train_mask_patches = gen_sub(y_train_mask,d,s)\n",
    "\n",
    "\n",
    "y_val_patches = gen_sub(y_val,d,s)\n",
    "y_val_mask_patches = gen_sub(y_val_mask,d,s)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# print first patch and check if LR is in accordance with HR\n",
    "fig, ax = plt.subplots(1,2, figsize=(10,10))\n",
    "ax[0].imshow(X_train_patches[0,:,:,0], cmap = 'gray')\n",
    "ax[1].imshow(y_train_patches[0,:,:,0], cmap = 'gray')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# free up memory\n",
    "del X_train, y_train, y_train_mask\n",
    "\n",
    "del X_val, y_val, y_val_mask"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.2 Clarity patches check"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# find patches indices with a lower percentage of clean pixels in train array\n",
    "patches_to_remove_train = [i for i,m in enumerate(y_train_mask_patches) if np.count_nonzero(m)/(HR_SIZE*HR_SIZE) < CLEAN_PATH_PX]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# find patches indices with a lower percentage of clean pixels in validation array\n",
    "patches_to_remove_val = [i for i,m in enumerate(y_val_mask_patches) if np.count_nonzero(m)/(HR_SIZE*HR_SIZE) < CLEAN_PATH_PX]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# print number of patches to be removed\n",
    "print(len(patches_to_remove_train))\n",
    "print(len(patches_to_remove_val))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# remove patches not clean\n",
    "X_train_patches = np.delete(X_train_patches,patches_to_remove_train,axis=0)\n",
    "y_train_patches =  np.delete(y_train_patches,patches_to_remove_train,axis=0)\n",
    "y_train_mask_patches =  np.delete(y_train_mask_patches,patches_to_remove_train,axis=0)\n",
    "\n",
    "X_val_patches = np.delete(X_val_patches,patches_to_remove_val,axis=0)\n",
    "y_val_patches =  np.delete(y_val_patches,patches_to_remove_val,axis=0)\n",
    "y_val_mask_patches =  np.delete(y_val_mask_patches,patches_to_remove_val,axis=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"network\"></a>\n",
    "# 3.0 Build the network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-03-06T14:18:41.261577Z",
     "start_time": "2020-03-06T14:18:38.172021Z"
    }
   },
   "outputs": [],
   "source": [
    "# build rams network\n",
    "rams_network = RAMS(scale=SCALE, filters=FILTERS, \n",
    "                 kernel_size=KERNEL_SIZE, channels=CHANNELS, r=R, N=N)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-03-06T14:18:41.365506Z",
     "start_time": "2020-03-06T14:18:41.297575Z"
    }
   },
   "outputs": [],
   "source": [
    "# print architecture structure\n",
    "rams_network.summary(line_length=120)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"train\"></a>\n",
    "# 4.0 Train the network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer_rams = Trainer(rams_network, band, HR_SIZE, name_net,\n",
    "                      loss=l1_loss,\n",
    "                      metric=psnr,\n",
    "                      optimizer=tf.keras.optimizers.Nadam(learning_rate=lr),\n",
    "                      checkpoint_dir=os.path.join(checkpoint_dir),\n",
    "                      log_dir=log_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer_rams.fit(X_train_patches,\n",
    "                [y_train_patches.astype('float32'), y_train_mask_patches], initial_epoch = 0,\n",
    "                batch_size=BATCH_SIZE, evaluate_every=400, data_aug = True, epochs=EPOCHS_N,\n",
    "                validation_data=(X_val_patches, [y_val_patches.astype('float32'), y_val_mask_patches])) "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Tensorflow2.0",
   "language": "python",
   "name": "tensorflow2.0"
  },
  "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.6.0"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": false,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": true
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}