test_RAMS_network.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Test 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 test the pre-trained or re-trained model (if you use the training notebook) using the weights stored in the 'ckpt' folder. The testing is performed with the validation set composed of all the scenes used to evaluate our network and all significant solutions presented in literature at the time of writing. Results can be found in the following table. Your aim is to beat the last row :)\n",
    "\n",
    "![table_results](media/probav_results_summary.png \"Logo Title Text 1\")\n",
    "\n",
    "**NB**: The last chapter of this notebook allows to make predictions with the original test set (ground truth are not public) of the Proba-V ESA challenge. After prediction and zip creation you can submit your file on the official post-mortem challenge, [here](https://kelvins.esa.int/proba-v-super-resolution-post-mortem/home/), and find out your score.\n",
    "\n",
    "**The notebook is divided in**:\n",
    "- 1.0 [Dataset Loading](#loading)\n",
    "- 2.0 [Load the Network](#network)\n",
    "- 3.0 [Test the network](#test)\n",
    "    - 3.1 Qualitative results\n",
    "    - 3.2 Compute RAMS cPSNR\n",
    "    - 3.3 Compute RAMS+ cPSNR\n",
    "    - 3.4 Compute RAMS SSIM\n",
    "    - 3.5 Computer RAMS+ SSIM\n",
    "- 4.0 [Predict Proba-V Test Set](#proba)\n",
    "    - 4.1 RAMS prediction\n",
    "    - 4.2 RAMS+ prediction\n",
    "    - 4.3 Submission zip creation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-01T16:31:48.339498Z",
     "start_time": "2020-10-01T16:31:48.132829Z"
    }
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "%reload_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-01T16:31:49.428001Z",
     "start_time": "2020-10-01T16:31:48.340430Z"
    }
   },
   "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.prediction import ensemble, unensemble, shuffle_last_axis, predict_tensor, predict_tensor_permute, savePredictions, savePredictionsPermut\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-10-01T16:31:49.460611Z",
     "start_time": "2020-10-01T16:31:49.429126Z"
    }
   },
   "outputs": [],
   "source": [
    "# gpu setting (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)\n",
    "tf.random.set_seed(42)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-01T16:35:34.451926Z",
     "start_time": "2020-10-01T16:35:34.431760Z"
    }
   },
   "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 = 128 # 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}' # weights path\n",
    "log_dir = 'logs' # tensorboard logs path\n",
    "submission_dir = 'submission' # submission dir\n",
    "name_zip = 'submission_RAMS.zip'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-01T16:35:35.625662Z",
     "start_time": "2020-10-01T16:35:35.591707Z"
    }
   },
   "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": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"loading\"></a>\n",
    "# 1.0 Dataset Loading"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-01T16:31:49.532375Z",
     "start_time": "2020-10-01T16:31:49.496607Z"
    }
   },
   "outputs": [],
   "source": [
    "# load validation \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": {
    "ExecuteTime": {
     "end_time": "2020-10-01T16:38:39.753944Z",
     "start_time": "2020-10-01T16:38:39.691011Z"
    }
   },
   "outputs": [],
   "source": [
    "# load ESA test set (no ground truth)\n",
    "X_test = np.load(os.path.join(PATH_DATASET, f'X_{band}_test.npy'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-01T16:38:41.923713Z",
     "start_time": "2020-10-01T16:38:41.898820Z"
    }
   },
   "outputs": [],
   "source": [
    "# print loaded dataset info\n",
    "print('X_val: ', X_val.shape)\n",
    "print('y_val: ', y_val.shape)\n",
    "print('y_val_mask: ', y_val_mask.shape)\n",
    "\n",
    "print('X_test: ', X_test.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"network\"></a>\n",
    "# 2.0 Load the Network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-01T16:35:43.341630Z",
     "start_time": "2020-10-01T16:35:40.864963Z"
    }
   },
   "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-10-01T16:35:43.541916Z",
     "start_time": "2020-10-01T16:35:43.427670Z"
    }
   },
   "outputs": [],
   "source": [
    "# load weights from checkpoint_dir\n",
    "checkpoint = tf.train.Checkpoint(step=tf.Variable(0),\n",
    "                                psnr=tf.Variable(1.0),\n",
    "                                model=rams_network)\n",
    "\n",
    "checkpoint.restore(tf.train.latest_checkpoint(checkpoint_dir))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"test\"></a>\n",
    "# 3.0 Test the Network"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3.1 Qualitative results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-01T16:35:53.048505Z",
     "start_time": "2020-10-01T16:35:52.201178Z"
    }
   },
   "outputs": [],
   "source": [
    "# print example images\n",
    "index = 15 # choose an image from validation set\n",
    "\n",
    "x_pred = predict_tensor(rams_network, X_val[index:index+1])\n",
    "\n",
    "fig, ax = plt.subplots(2,2, figsize=(15,15))\n",
    "ax[0,0].imshow(X_val[index,:,:,0], cmap = 'gray')\n",
    "ax[0,0].set_title('LR')\n",
    "ax[0,1].imshow(bicubic(np.mean(X_val[index:index+1], axis=-1)[...,None])[0,:,:,0], cmap ='gray')\n",
    "ax[0,1].set_title('Bicubic')\n",
    "ax[1,0].imshow(x_pred[0,:,:,0], cmap ='gray')\n",
    "ax[1,0].set_title('Prediction')\n",
    "ax[1,1].imshow(y_val[index,:,:,0], cmap = 'gray')\n",
    "ax[1,1].set_title('HR')\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3.2 Compute RAMS cPSNR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-01T16:32:32.628640Z",
     "start_time": "2020-10-01T16:31:55.181620Z"
    }
   },
   "outputs": [],
   "source": [
    "# compute cPSNR with trained network\n",
    "psnr_scores = []\n",
    "\n",
    "for index in tqdm(range(X_val.shape[0])):\n",
    "    x_pred = predict_tensor(rams_network, X_val[index:index+1])\n",
    "    psnr_scores.append(psnr(y_val[index:index+1,:,:], x_pred, y_val_mask[index:index+1,:,:], HR_SIZE).numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-01T16:32:32.646385Z",
     "start_time": "2020-10-01T16:32:32.629638Z"
    }
   },
   "outputs": [],
   "source": [
    "# print validation network cPSNR\n",
    "print(f'PSNR Validation Network: {np.mean(psnr_scores)}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-01T16:32:51.498802Z",
     "start_time": "2020-10-01T16:32:32.647104Z"
    }
   },
   "outputs": [],
   "source": [
    "# compute cPSNR with bicubic interpolation\n",
    "psnr_scores_bicubic = []\n",
    "\n",
    "for index in tqdm(range(X_val.shape[0])):\n",
    "    x_pred = bicubic(np.mean(X_val[index:index+1], axis=-1)[...,None])\n",
    "    psnr_scores_bicubic.append(psnr(y_val[index:index+1,:,:], x_pred, y_val_mask[index:index+1,:,:], HR_SIZE).numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-01T16:32:51.516081Z",
     "start_time": "2020-10-01T16:32:51.499656Z"
    }
   },
   "outputs": [],
   "source": [
    "# print validation bicubic cPSNR\n",
    "print(f'PSNR Validation Bicubic: {np.mean(psnr_scores_bicubic)}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-01T16:32:51.663515Z",
     "start_time": "2020-10-01T16:32:51.516913Z"
    }
   },
   "outputs": [],
   "source": [
    "# show bicubic vs network comparison\n",
    "plt.rc('font', size=25)          # controls default text sizes\n",
    "fig,ax = plt.subplots(figsize=(10,10))\n",
    "\n",
    "ax.plot(np.arange(30,70,0.5),np.arange(30,70,0.5),'black',linewidth=0.5)\n",
    "ax.fill_between(np.arange(20,70,0.5), 0, np.arange(20,70,0.5), color = (0/255, 0/255, 230/255, 0.3))\n",
    "ax.fill_between(np.arange(20,70,0.5), np.arange(20,70,0.5), 70, color = (0/255, 0/255, 230/255, 0.1))\n",
    "\n",
    "ax.plot(psnr_scores_bicubic, psnr_scores,'x',color='black')\n",
    "\n",
    "plt.xticks(np.arange(30, 65, step=10))\n",
    "plt.yticks(np.arange(30, 65, step=10))\n",
    "plt.ylabel(\"mPSNR RAMS\")\n",
    "plt.xlabel(\"mPSNR Bicubic\")\n",
    "ax.axis((25,65,25,65))\n",
    "ax.set_aspect('equal')\n",
    "ax.grid(color = 'gray',linestyle='dotted')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3.3 Compute RAMS+ cPSNR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-01T16:38:50.588057Z",
     "start_time": "2020-10-01T16:38:45.965669Z"
    }
   },
   "outputs": [],
   "source": [
    "# compute Ensemble PSNR prediction\n",
    "psnr_scores_plus = []\n",
    "n_permut = 20 # number of permutations\n",
    "\n",
    "for index in tqdm(range(X_val.shape[0])):\n",
    "    x_pred = predict_tensor_permute(rams_network, X_val[index],n_ens=n_permut)\n",
    "    psnr_scores_plus.append(psnr(y_val[index:index+1,:,:], x_pred, y_val_mask[index:index+1,:,:], HR_SIZE).numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-01T16:35:18.137411Z",
     "start_time": "2020-10-01T16:31:48.171Z"
    }
   },
   "outputs": [],
   "source": [
    "# print validation Ensamble cPSNR\n",
    "print(f'PSNR Validation Ensamble: {np.mean(psnr_scores_plus)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3.4 Compute RAMS SSIM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-01T16:36:58.569309Z",
     "start_time": "2020-10-01T16:36:56.182755Z"
    }
   },
   "outputs": [],
   "source": [
    "# compute SSIM with trained network\n",
    "ssim_scores = []\n",
    "\n",
    "for index in tqdm(range(X_val.shape[0])):\n",
    "    x_pred = predict_tensor(rams_network, X_val[index:index+1])\n",
    "    ssim_scores.append(ssim(y_val[index:index+1,:,:], x_pred, y_val_mask[index:index+1,:,:], HR_SIZE).numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-01T16:35:18.138701Z",
     "start_time": "2020-10-01T16:31:48.176Z"
    }
   },
   "outputs": [],
   "source": [
    "# print validation network SSIM\n",
    "print(f'SSIM Validation Network: {np.mean(ssim_scores)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3.5 Computer RAMS+ SSIM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-01T16:35:18.139307Z",
     "start_time": "2020-10-01T16:31:48.179Z"
    }
   },
   "outputs": [],
   "source": [
    "# compute Ensemble SSIM prediction\n",
    "ssim_scores_ens = []\n",
    "n_permut = 20 # number of permutations\n",
    "\n",
    "for index in tqdm(range(X_val.shape[0])):\n",
    "    x_pred = predict_tensor_permute(rams_network, X_val[index],n_ens=n_permut)\n",
    "    ssim_scores_ens.append(ssim(y_val[index:index+1,:,:], x_pred, y_val_mask[index:index+1,:,:], HR_SIZE).numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-01T16:35:18.139933Z",
     "start_time": "2020-10-01T16:31:48.181Z"
    }
   },
   "outputs": [],
   "source": [
    "# print validation Ensamble SSIM\n",
    "print(f'PSNR Validation: {np.mean(ssim_scores_ens)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"proba\"></a>\n",
    "# 4.0 Predict Proba-V Test Set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-01T16:36:04.975830Z",
     "start_time": "2020-10-01T16:36:04.952158Z"
    }
   },
   "outputs": [],
   "source": [
    "# create submission folder\n",
    "if not os.path.exists(submission_dir):\n",
    "    os.mkdir(submission_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4.1 RAMS prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-01T16:37:22.317278Z",
     "start_time": "2020-10-01T16:37:22.293368Z"
    }
   },
   "outputs": [],
   "source": [
    "X_val.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-01T16:37:15.246901Z",
     "start_time": "2020-10-01T16:37:15.221858Z"
    }
   },
   "outputs": [],
   "source": [
    "X_test.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-01T16:38:58.851952Z",
     "start_time": "2020-10-01T16:38:56.061194Z"
    }
   },
   "outputs": [],
   "source": [
    "# predict proba-v test set with RAMS\n",
    "X_preds = []\n",
    "\n",
    "for index in tqdm(range(X_test.shape[0])):\n",
    "    X_preds.append(predict_tensor(rams_network, X_test[index:index+1]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-01T16:35:18.141794Z",
     "start_time": "2020-10-01T16:31:48.188Z"
    }
   },
   "outputs": [],
   "source": [
    "savePredictions(X_preds, band, submission_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4.2 RAMS+ prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-01T16:35:18.142520Z",
     "start_time": "2020-10-01T16:31:48.190Z"
    }
   },
   "outputs": [],
   "source": [
    "# predict proba-v test set with RAMS+\n",
    "X_preds = []\n",
    "n_permut = 20 # number of permutations\n",
    "\n",
    "for index in tqdm(range(X_test.shape[0])):\n",
    "    X_preds.append(predict_tensor_permute(rams_network, X_test[index], n_ens=n_permut))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-01T16:35:18.143130Z",
     "start_time": "2020-10-01T16:31:48.193Z"
    }
   },
   "outputs": [],
   "source": [
    "# save predictions in submission_dir\n",
    "savePredictionsPermut(X_preds, band, submission_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4.3 Submission zip creation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-10-01T16:35:18.143709Z",
     "start_time": "2020-10-01T16:31:48.196Z"
    }
   },
   "outputs": [],
   "source": [
    "# zip creation\n",
    "zf = ZipFile(name_zip, mode='w')\n",
    "with tqdm(total=290, desc=\"Zipping images\") as pbar:\n",
    "    for i, img in enumerate(sorted(os.listdir(submission_dir))):\n",
    "        zf.write(os.path.join(submission_dir, img), arcname=img)\n",
    "        pbar.update(1)\n",
    "zf.close()\n",
    "print('Done!')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "tensorflow",
   "language": "python",
   "name": "tensorflow"
  },
  "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.9"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "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": false
  },
  "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": 2
}