notebooks/dlmbl2022/exercise.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Virtual staining by image translation\n",
    "***\n",
    "\n",
    "This notebook illustrates image translation with deep convolutional neural networks (CNN). We frame the image translation as a regression problem solved using a residual U-Net model. The notebook demonstrates how to translate quantitative phase images of mouse kidney tissue to the fluorescence images of nuclei using data and model reported in our paper(https://doi.org/10.7554/eLife.55502).\n",
    "<div>\n",
    "    <img src=\"virtualstaining_summary.png\"  width=\"950\">\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "##  Overview <a class=\"anchor\" id=\"overview\"></a>\n",
    "\n",
    "U-Net model consists of an encoder (downsampling) part and a decoder (upsampling) part. The U-Net model is immensely popular for many image analysis tasks. One of the key design feature of the U-Net architecture is the skip connections between the encoder and decoder layers, which allows the model to learn patterns at multiple spatial resolutions. U-Nets were orginally designed for image segmentation (https://link.springer.com/chapter/10.1007/978-3-319-24574-4_28). Using U-Nets for image translation (a regression problem) needs a few tweaks, such as: \n",
    "* changing the loss function, \n",
    "* changing the final activation layer, and \n",
    "* data normalization (more on that later).\n",
    "\n",
    "Below is an overview of the 2D and 2.5D residual U-Net model architecture. The main differences from the original U-Net architecture are: \n",
    "1. We add short skip connection within each block (residual block) in addition to the long skip connections in the original U-Net. \n",
    "2. The long skip connections in the 2.5D model project the data along depth dimension - it uses short stacks (5-7 slices) as input to compute a 2D output. This architecture provides better translation accuracy than 2D model, because the model learns the pattern of blur along depth. \n",
    "3. The 2.5D model can be used to virtually stain 3D stacks by sliding the input window over the data along $x$, $y$, and $z$ dimensions.\n",
    "\n",
    "<div>\n",
    "    <img src=\"supp_modelarch_RGB.png\"  width=\"950\">\n",
    "</div>\n",
    "The 2D model translates slice$\\rightarrow$slice, whereas 2.5D model translates stack$\\rightarrow$slice. \n",
    "\n",
    "We'll be using the architecture similar to above, but with same convolution instead of valid convolution.  Same convolution operation pads the input image so that the output image has the same size as the input image after convolution.\n",
    "\n",
    "***\n",
    "The first half of the exercise will focus on predicting nuclei from label-free measurements using different flavors of 2D UNets in [microDL](https://github.com/mehta-lab/microDL) in the jupyter notebook. \n",
    "In the second half, you can fine tune the model or explore translation between different types of modalities. If you are comfortable using CLI, you can potentially train 2-3 small models in parallel on your GPU using config files.\n",
    "***"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set path and import libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "\n",
    "TOP_DIR = os.path.expanduser('.') # Top directory in which microDL library, data, patches, and trained models are saved.'.' means current directory. \n",
    "INPUT_DIR = os.path.join(TOP_DIR, 'data') # directory where raw data is saved.\n",
    "TILE_DIR = os.path.join(TOP_DIR, 'phase2dna_microdl_patches') # directory where patches will be saved.\n",
    "MODEL_DIR = os.path.join(TOP_DIR, 'microdl_model') # directory where model will be saved.\n",
    "module_path = os.path.join(TOP_DIR, 'microDL') #directory of microDL module.\n",
    "\n",
    "# make sure that microDL library can be seen by this interpreter.\n",
    "if module_path not in sys.path:\n",
    "    sys.path.append(module_path)\n",
    "print(sys.path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# autoreload modules if they have changed.\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import glob\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import importlib\n",
    "import pandas as pd\n",
    "import pprint\n",
    "\n",
    "# Supress warnings related to deprecation of tensorflow 1.x    \n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\", message=r\"Passing\", category=FutureWarning)\n",
    "warnings.filterwarnings(\"ignore\", message=r\"DEPRECATED\", category=UserWarning)\n",
    "\n",
    "import tensorflow as tf\n",
    "tf.logging.set_verbosity(tf.logging.ERROR)\n",
    "\n",
    "import micro_dl.utils.meta_utils as meta_utils\n",
    "import micro_dl.cli.train_script as train\n",
    "import micro_dl.cli.preprocess_script as preprocess\n",
    "import micro_dl.utils.aux_utils as aux_utils\n",
    "import micro_dl.utils.image_utils as im_utils\n",
    "import micro_dl.utils.masks as mask_utils\n",
    "import micro_dl.utils.normalize as norm_utils\n",
    "import micro_dl.utils.train_utils as train_utils\n",
    "import micro_dl.inference.image_inference as image_inf\n",
    "import micro_dl.cli.metrics_script as metrics\n",
    "    \n",
    "# Setup pretty printing\n",
    "pp = pprint.PrettyPrinter(indent=4)\n",
    "\n",
    "\n",
    "font = {'family' : 'DejaVu Sans',\n",
    "        'size'   : 20}\n",
    "\n",
    "matplotlib.rc('font', **font)\n",
    "%matplotlib inline\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Explore data and metadata\n",
    "\n",
    "Data examination and curation is a *very* important part of training accurate and useful models. ML researchers typically spend half or more of their time curating the right dataset for the problem at hand. In this case, we will start with a small dataset to quickly iterate on the parameters of a virtual staining model, before refining the accuracy with microDL."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Explore input and target images"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The goal of this exercise is to translate 2D slices of phase or retardance image into 2D slices of fluorescene images of F-actin and DNA. The specimen is a mouse kidney tissue section. Let's start by looking at some example images."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-info\"><h4>\n",
    "    TODO 1: Explore the label-free and fluorescence images of the same FOV."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "# hint: paths look like this.\n",
    "position = 5\n",
    "actin_path = os.path.join(INPUT_DIR, f'img_568_t000_p{position:03}_z010.tif') #F-actin was imaged with Alexa Fluor 568 using 468nm excitation wavelength.\n",
    "dna_path=os.path.join(INPUT_DIR, f'img_405_t000_p{position:03}_z010.tif') # DNA was imaged with Hoechst using 405nm excitation wavelength.\n",
    "phase_path=os.path.join(INPUT_DIR,f'img_phase_t000_p{position:03}_z010.tif') # the phase and polarization were imaged using 530nm wavelength.\n",
    "ret_path=os.path.join(INPUT_DIR,f'img_Retardance_t000_p{position:03}_z010.tif') # the phase and polarization were imaged using 530nm wavelength.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Metadata <a class=\"anchor\" id=\"metadata\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "microDL uses CSV file to track the metadata of the images. We parse metadata data using a widely used [pandas](https://pandas.pydata.org/) library.  We'll generate the metadata and take a look:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "meta_utils.frames_meta_generator(INPUT_DIR)\n",
    "meta_utils.ints_meta_generator(INPUT_DIR,\n",
    "                               num_workers=8,\n",
    "                               )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "frames_meta = pd.read_csv(os.path.join(INPUT_DIR, 'frames_meta.csv'), index_col=0)\n",
    "frames_meta"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Each channel in our data is assgined a unique channel ID. We'll reference the channels by their IDs in the config files. Let's check the channel IDs. `pandas` frames can be indexed into like numpy arrays where index can be the headers of a column. See more examples [here](https://pandas.pydata.org/docs/user_guide/indexing.html).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "chan_id_df = frames_meta[['channel_idx', 'channel_name']].drop_duplicates()\n",
    "chan_id_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Find position ids \n",
    "pos_id_df = frames_meta[['pos_idx']].drop_duplicates()\n",
    "pos_id_df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Checkpoint 1\n",
    "\n",
    "<div class=\"alert alert-block alert-success\"><h4>\n",
    "    \n",
    "Let us know by slack if you have been able to load the data and generate metadata table."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Preprocessing <a class=\"anchor\" id=\"preproc\"></a>\n",
    "\n",
    "We will  be using three main modules of microDL for image translation: preprocessing, training and inference (see imports above).\n",
    "\n",
    "The first step is the preprocessing. The key step in the preprocessing is tiling the images. The original image size (2048 x 2048 pixels) is too large to be able to fit into memory and similar structures reappear across the images. So we'd be much better off splitting the images into smalle pieces (patches). A design consideration for convolution neural networks is the [receptive field](https://medium.com/mlreview/a-guide-to-receptive-field-arithmetic-for-convolutional-neural-networks-e0f514068807), which is the region of the input space a given feature is looking at. Given the size of our structures of interest and our network architecture, we use 256 x 256 pixel tiles. Also, training on smaller tiles allow us to use bigger batch size so the training converges faster.\n",
    "\n",
    "We have additional options like resizing, flatfield correction, and creating masks. The data we're working with is already background corrected so we can safely skip that part. We also don't need to resize the images. We would however like to create masks based on our target data. The reason for that is that we would like to avoid training on background only, that is empty or near empty tiles. A threshold we found reasonable is to make sure that 25% or more pixels of our tiles contain signal.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### masks\n",
    "\n",
    "We would like to generate binary masks based on our target channel. In microDL we have two types of global binary thresholding methods builtin, [Otsu](https://en.wikipedia.org/wiki/Otsu%27s_method) and [unimodal](https://users.cs.cf.ac.uk/Paul.Rosin/resources/papers/unimodal2.pdf) (or Rosin) thresholding.\n",
    "Let's load an image, generate masks and plot them side by side."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load image\n",
    "position = 5\n",
    "im_path = os.path.join(INPUT_DIR, f'img_405_t000_p{position:03}_z010.tif')\n",
    "im = im_utils.read_image(im_path)\n",
    "# Clip top and bottom 1% of histogram for better visualization\n",
    "im = norm_utils.hist_clipping(im, 1, 99)\n",
    "\n",
    "mask_otsu = mask_utils.create_otsu_mask(im)\n",
    "mask_rosin = mask_utils.create_unimodal_mask(im)\n",
    "\n",
    "fig, ax = plt.subplots(1, 3)\n",
    "fig.set_size_inches(20, 15)\n",
    "ax = ax.flatten()\n",
    "ax[0].imshow(im, cmap='gray')\n",
    "ax[0].set_title('Fluorescence',fontsize=20)\n",
    "ax[1].imshow(mask_otsu, cmap='gray')\n",
    "ax[1].set_title('Otsu thresholding',fontsize=20)\n",
    "ax[2].imshow(mask_rosin, cmap='gray')\n",
    "ax[2].set_title('unimodal thresholding',fontsize=20)\n",
    "for a in ax: a.axis('off')\n",
    "\n",
    "plt.show()\n",
    "plt.close('all')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### configuration"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For analyzing large amounts of data and to keep track of the parameters, we run microDL through command line interfaces (CLIs) using configuration files. For each module we specify the path to a config file as a command line argument. Since we're using Jupyter Notebook for this tutorial we will instead load the preprocessing config and update it. Look for `...` to fill in the blanks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "config_path = os.path.join(\n",
    "    module_path,\n",
    "    'config_files/config_preprocess.yml',\n",
    ")\n",
    "preproc_config = aux_utils.read_config(config_path)\n",
    "\n",
    "# We're not doing resizing so let's remove that key\n",
    "if 'resize' in preproc_config:\n",
    "    preproc_config.pop('resize')\n",
    "# We're not doing flatfield correction either\n",
    "if 'flat_field' in preproc_config:\n",
    "    preproc_config.pop('flat_field')\n",
    "    \n",
    "# We need to change INPUT_DIR to point to where our image data is located\n",
    "preproc_config['input_dir'] = INPUT_DIR\n",
    "\n",
    "# And where we want to store our preprocessed data\n",
    "preproc_config['output_dir'] = TILE_DIR\n",
    "\n",
    "# Set which channels we want to preprocess and if they should be normalized\n",
    "preproc_config['channel_ids'] = [0, 3]\n",
    "preproc_config['normalize']['normalize_channels'] = [True, True]\n",
    "preproc_config['tile']['depths'] = [1, 1]\n",
    "\n",
    "# Set the channels used for generating masks\n",
    "preproc_config['masks']['channels'] = 0\n",
    "\n",
    "# Set the number of workers to the number of available cores\n",
    "preproc_config['num_workers'] = 8"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-info\"><h4>\n",
    "    TODO 2: Choose a masking strategy for this exercise that you think will result in the optimal balance between foreground and background pixels. This begs a question: for regression problems, which pixels are foreground?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you can see, the Otsu threshold captures only the very bright structures and misses dim structures. Rosin thresholding does a better job picking up these dim structures."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "preproc_config['masks']['mask_type'] = ... # TODO 2: your choice of thresholding algorithm here.\n",
    "preproc_config['pos_ids'] = [3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 40, 41, 42]\n",
    "pp.pprint(preproc_config)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### generate patches/tiles\n",
    "Now it's time to run the preprocessing: the runtime should be around 5 minute. \n",
    "***\n",
    "Notes:\n",
    "1. If you get an error about \"Don't specify a mask_dir\", try reloading the config by running the last two blocks.\n",
    "2. If above step takes too long due to network or data i/o bottlenecks, you can copy the patches from this location to the TILE_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "preproc_config, runtime = preprocess.pre_process(preproc_config);\n",
    "print(\"Preprocessing took {} seconds\".format(runtime))\n",
    "# Save the final config and run time\n",
    "preprocess.save_config(preproc_config, runtime)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "You can navigate to your output_dir and take a look at what was generated. You will find a mask_dir containing masks, a tile_dir containing tiles and JSON file containing the all the information that was used to generated the preprocessed data. Let's take a look at a few tiles. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### visualize 10 random patches"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tile_dir = preproc_config['tile']['tile_dir']\n",
    "# tile_dir =os.path.join(os.path.expanduser('~'), '04_image_translation_data/tmp/tile_small/tiles_256-256_step_128-128' )\n",
    "print(tile_dir)\n",
    "frames_meta = pd.read_csv(os.path.join(tile_dir, 'frames_meta.csv'))\n",
    "dna_tiles = frames_meta[frames_meta.file_name.str.match('im_c000_*')]\n",
    "subset = np.random.choice(dna_tiles.shape[0], 10, replace=False)\n",
    "\n",
    "fig, ax = plt.subplots(3, 10)\n",
    "fig.set_size_inches(20, 10)\n",
    "# ax = ax.flatten()\n",
    "for i, row  in enumerate(subset):\n",
    "    dna_tile_path = os.path.join(tile_dir, frames_meta.iloc[row].file_name)\n",
    "    phase_tile_path = dna_tile_path.replace('c000','c003')\n",
    "    mask_tile_path = dna_tile_path.replace('c000','c004')\n",
    "    \n",
    "    dna = im_utils.read_image(dna_tile_path)\n",
    "    phase = im_utils.read_image(phase_tile_path)\n",
    "    mask = im_utils.read_image(mask_tile_path)\n",
    "    ax[0,i].imshow(np.squeeze(dna), cmap='gray'); ax[0,i].axis('off')\n",
    "    ax[1,i].imshow(np.squeeze(phase), cmap='gray'); ax[1,i].axis('off')\n",
    "    ax[2,i].imshow(np.squeeze(mask), cmap='gray'); ax[2,i].axis('off')\n",
    "  \n",
    "plt.show()\n",
    "plt.close('all')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Checkpoint 2\n",
    "\n",
    "<div class=\"alert alert-block alert-success\"><h4>\n",
    "    \n",
    "Let us know by slack if you have successfully generated tiles for image translation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. 2D virtual staining (slice$\\rightarrow$slice)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Configure training\n",
    "Now that we've preprocessed our data we're ready to train. During this exercise, we will use phase or retardance image to predict two different fluorescent channels, actin and nuclei.\n",
    "In our dataset, the channel names and indices are the following:\n",
    "\n",
    "* Retardance: channel name Retardance, index 2\n",
    "* Phase: channel name phase, index 3\n",
    "* Actin:channel name 568, index 1\n",
    "* Nuclei: channel name 405, index 0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We specify the network architecture and training parameters using another config file. Let's load a base 2D training config file and take a look.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config_path = os.path.join(\n",
    "    module_path,\n",
    "    'config_files/config_train.yml',\n",
    ")\n",
    "train_config = aux_utils.read_config(config_path)\n",
    "# Set the data directory to the directory we generated during preprocessing\n",
    "train_config['dataset']['data_dir'] = os.path.join(TILE_DIR, 'tiles_256-256_step_128-128')\n",
    "\n",
    "# We also need to specify where we want to store our model and all related data generated by training\n",
    "# This directory will be created if it doesn't already exist\n",
    "train_config['trainer']['model_dir'] = MODEL_DIR\n",
    "\n",
    "# Predict dna (channel 1) from phase (channel 3)\n",
    "train_config['dataset']['input_channels'] = [3]\n",
    "train_config['dataset']['target_channels'] = [0]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Train the model!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There are many parameters of a model architecture and training process that a DL practitioner optimizes by hand - these are broadly termed hyper-paramters. They are distinct from parameters of the layers of the neural network learned from the data. Following are key hyper-parameters for virtual staining and good defaults.\n",
    "\n",
    "#### Final activation layer\n",
    "Virtual staining is a regression task,and therefore the final activation layer needs to be linear. For a binary segmentation, we use a sigmoid to move output towards either zero or one. For regression, we would like the prediction to match the dynamic range of the target.\n",
    "\n",
    "##### Loss\n",
    "Common choices for regression are the mean squared error (MSE) and the mean absolute error (MAE) between the target image y and the estimated image y':\n",
    "\\begin{equation*}\n",
    "MSE = \\sum_{p} (y_p - y_p')^2,\n",
    "\\end{equation*}\n",
    "\n",
    "\\begin{equation*}\n",
    "MAE = \\sum_{p} |y_p - y_p'|\n",
    "\\end{equation*}\n",
    "for each pixel index p.\n",
    "\n",
    "You can try both by changing `train_config['trainer']['loss']`. The names in microDL for losses are mse_loss and mae_loss, respectively. You can also try other custom losses by using the loss function names [here](https://github.com/czbiohub/microDL/blob/master/micro_dl/train/losses.py), or any standard [Keras loss function](https://keras.io/losses/) by specifying the loss function name defined by Keras. \n",
    "\n",
    "##### Optimizer\n",
    "Adam is a good default optimizer. You can read more about different deep learning optimizers [here](http://ruder.io/optimizing-gradient-descent/), and you can change the optimizer you'd like to use in your training by changing the variable `train_config['trainer']['optimizer']['name']` to any of the Keras optimizers listed [here](https://keras.io/optimizers/).\n",
    "\n",
    "##### Learning Rate\n",
    "If the learning rate is too small your training might take a very long time to converge, and if it's too big it might not converge at all. It's worth trying some different values and see what happens with convergence.\n",
    "\n",
    "##### Dropout\n",
    "Since we're working with a very small dataset in exploratory training, chances are that your network will overfit to your training data. It's worth exploring `train_config['network']['dropout']` and to see if increasing those variables can reduce overfitting.\n",
    "\n",
    "##### Number of filters\n",
    "The number of filters in each layer of the model controls the model capacity. This parameter is `train_config['network']['num_filters_per_block']`. Too large model capacity can lead to overfitting and not necesssarily better results.\n",
    "\n",
    "##### Augmentation\n",
    "The data is flipped and rotated randomly to diversify the training set and mitigate overfitting.\n",
    "\n",
    "##### Other?\n",
    "If you have extra time or are curious about the other variables, feel free to explore."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-info\"><h4>\n",
    "    TODO 3: Pick a loss function and other hyper-parameters you think are right for the regression problem. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Enable data augmentation.\n",
    "train_config['dataset']['augmentation'] = True\n",
    "\n",
    "train_config['trainer']['loss'] = ... # TODO 3: your choice of loss function here.\n",
    "# Set maximum number of epochs to 7 so we can explore difference parameters quickly\n",
    "train_config['trainer']['max_epochs'] = 7\n",
    "# Batch size of 64 just about occupies 16GB RAM.\n",
    "train_config['trainer']['batch_size'] = 64\n",
    "\n",
    "# Set number of filters per block. The size of this vector sets the depth of U-Net.\n",
    "train_config['network']['num_filters_per_block'] = [16, 32, 48, 64, 80]\n",
    "# You can also update other training hyper-parameters here.\n",
    "pp.pprint(train_config)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# First, specify the gpu that you want to run training on, or \"None\" to select the gpu with most memory available\n",
    "gpu_id = 0\n",
    "gpu_id, gpu_mem_frac = train_utils.select_gpu(gpu_ids=gpu_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Try restarting the notebook kernel if the gpu memory is occupied and you run into errors about \"can't creat training session\".\n",
    "Training 7 epochs should take no more than 10 minutes if you're on a GPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train.run_action(\n",
    "    action='train',\n",
    "    config=train_config,\n",
    "    gpu_ids=gpu_id,\n",
    "    gpu_mem_frac=gpu_mem_frac,\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We've completed our first training. Let's take a look at what happened during training by opening a history log file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "history = pd.read_csv(os.path.join(MODEL_DIR, 'history.csv'))\n",
    "history"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot training vs. validation loss\n",
    "fig, ax = plt.subplots()\n",
    "fig.set_size_inches(12, 9)\n",
    "ax.plot(history['epoch'], history['loss'], 'r')\n",
    "ax.plot(history['epoch'], history['val_loss'], 'b')\n",
    "ax.set_xlabel('Epoch')\n",
    "ax.set_ylabel('Loss')\n",
    "plt.legend(['train', 'validation'])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "toc-hr-collapsed": true
   },
   "source": [
    "### Checkpoint 3\n",
    "\n",
    "<div class=\"alert alert-block alert-success\"><h4>\n",
    "Please let us know by slack whether your training converged.\n",
    "    \n",
    "**Does it look like the model is overfitting? How can you tell?**\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## 4. Predictions on test set <a class=\"anchor\" id=\"predict\"></a>\n",
    "\n",
    "We'd also like to see how well the model performs predictions. For that we will have to run inference on our test dataset. We will run model inference on the full size 2048 X 2048 image instead of on tiles in training. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-info\"><h4>\n",
    "\n",
    "TODO 4: Why can we run the model inference with input size different from when used during training? And what are the benefits of doing that?  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "... TODO 4: your response"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config_path = os.path.join(\n",
    "    module_path,\n",
    "    'config_files/config_inference.yml',\n",
    ")\n",
    "inf_config = aux_utils.read_config(config_path)\n",
    "inf_config['image_dir'] = INPUT_DIR \n",
    "inf_config['preprocess_dir'] = TILE_DIR\n",
    "inf_config['model_dir'] = MODEL_DIR\n",
    "inf_config['dataset']['input_channels'] = [3]\n",
    "inf_config['dataset']['target_channels'] = [0]\n",
    "inf_config['metrics']['metrics_orientations'] = ['xy']\n",
    "pp.pprint(inf_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "split = aux_utils.read_config(os.path.join(MODEL_DIR , 'split_samples.json'))\n",
    "split"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "inference_inst = image_inf.ImagePredictor(\n",
    "                train_config=train_config,\n",
    "                inference_config=inf_config,\n",
    "                preprocess_config=preproc_config,\n",
    "                gpu_id=gpu_id,\n",
    "                gpu_mem_frac=gpu_mem_frac,\n",
    "            )\n",
    "inference_inst.run_prediction()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There should be a new subdirectory created in the model directory with the predictions and the metrics. Use glob to see what files were generated during inference."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot an example of input, target and prediction side by side."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pred_dir = os.path.join(MODEL_DIR, 'predictions')\n",
    "position = split['test'][0]\n",
    "z = 12\n",
    "fig, ax = plt.subplots(1, 3)\n",
    "fig.set_size_inches(20, 15)\n",
    "ax = ax.flatten()\n",
    "for a in ax: a.axis('off')\n",
    "\n",
    "im_path = os.path.join(INPUT_DIR, f'img_phase_t000_p{position:03}_z{z:03}.tif')\n",
    "im = im_utils.read_image(im_path)\n",
    "im = norm_utils.hist_clipping(im, 1, 99)\n",
    "ax[0].imshow(im, cmap='gray'); ax[0].set_title('Input: Phase', fontsize=20)\n",
    "im_path = os.path.join(INPUT_DIR, f'img_405_t000_p{position:03}_z{z:03}.tif')\n",
    "im = im_utils.read_image(im_path)\n",
    "im = norm_utils.hist_clipping(im, 1, 99)\n",
    "ax[1].imshow(im, cmap='gray'); ax[1].set_title('Target: DNA', fontsize=20)\n",
    "im_path = os.path.join(pred_dir, f'img_405_t000_p{position:03}_z{z:03}.tif')\n",
    "im = im_utils.read_image(im_path)\n",
    "im = norm_utils.hist_clipping(im, 1, 99)\n",
    "ax[2].imshow(im, cmap='gray'); ax[2].set_title('Prediction of DNA', fontsize=20)\n",
    "plt.show()\n",
    "plt.close('all')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The prediction is blurrier than the target. One reason you can't exactly mimic the nuclei image is that the input channel carries partial information about the structure and the random noise cannot be learned.\n",
    "\n",
    "Another reason for blurry prediction is that loss functions like MSE or MAE tend to generate blurrier prediction because these loss functions generate an \"average\" prediction when there are multiple possible predictions.  \n",
    "\n",
    "Also, we've here used a very limited amount of data. To get higher correlation we will need to include much more data and run training overnight.\n",
    "\n",
    "Speaking of correlation, let's open the inference meta file and inspect the metrics comparing predictions and targets.\n",
    "\n",
    "Each metric reports on different aspects of the fit between predictions (P) and target (T).  \n",
    "* $R^2$ is the coefficient of determination = $1- \\frac{MSE(P,T)}{\\sigma^2(T)}$. It is sensitive to difference in the mean values.\n",
    "* The Pearson correlation coefficient is\n",
    "$$ r(T,P) = \\frac{\\sigma_{TP}}{\\sigma_T \\sigma_P} $$ It is insensitive to the mean value, and reports if two variables vary linearly.  \n",
    "* SSIM is a multi-component metric [developed to represent human vision](http://www.cns.nyu.edu/pub/lcv/wang03-reprint.pdf). It is a single metric that reports on similarity of luminance (local mean), contrast (local variance), and structure (local pearson correlation) between two images.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "metrics_meta = pd.read_csv(os.path.join(pred_dir, 'metrics_xy.csv'))\n",
    "metrics_meta.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Get the median correlation of all rows in the inference csv file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Median Pearson correlation: {:.2f}\".format(metrics_meta['corr'].median()))\n",
    "print(\"Median SSIM: {:.2f}\".format(metrics_meta['ssim'].median()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "toc-hr-collapsed": true
   },
   "source": [
    "### Checkpoint 4\n",
    "<div class=\"alert alert-block alert-success\"><h4>\n",
    "\n",
    "**Post the median metrics achieved with your model on slack along with key choices of hyper-parameters.** \n",
    "\n",
    "**Which aspects of data and the model lead to  less-than-perfect match?** "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Model tuning <a class=\"anchor\" id=\"tuning\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now it's time for you to experiment. You can try modeling a different channel (use retardance channel to predict F-actin) or play around with different settings in the train_config and rerun the training. What do you think will help improve the model?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Bonus 1: Training with different input-output pairs.\n",
    "\n",
    "With preprocessed tiles, you are setup to explore how accurately you can predict one channel from the other. Some questions you may want to explore:\n",
    "\n",
    "* How accurately can you predict the phase images from images of nuclei? Is the accuracy of prediction symmetric?\n",
    "\n",
    "* Which label-free channel (phase or retardance) provides more accurate prediction of F-actin?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##  Bonus 2: Training on larger dataset (slice$\\rightarrow$slice) overnight.\n",
    "\n",
    "Hopefully above exploration has led you to set of parameters to predict nuclei with decent accuracy. You can now set up a model to train on larger dataset (30 FOVs) and evaluate if model accuracy improves as a result when you comeback. \n",
    "\n",
    "We will need change the position ids in the pre-process config to have 30 FOVs and re-run preprocessing. The position ids of 30 FOVs are:\n",
    "\n",
    "[3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99]\n",
    "\n",
    "You can pick training/model parameters based on your parameter exploration from Session 1. A few tips for picking training parameters for overnight training:\n",
    "- **Make sure you write preprocessed data and model trained on this large set to new folders.**\n",
    "- Increase maximum number of epochs and early stopping patience to at least 200 and 10 so the training will run longer\n",
    "- Increase the number of filters in the model to increase the model capacity. You might need to use smaller batch size so the model can fit into the GPU memory You might want to add some dropout as well to avoid overfitting if you increase the number of filters \n",
    "- Use lower learning rate. We used higher learning rate to make training converge faster\n",
    "- Compare the mean and standard deviation of test metrics.\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Raw Cell Format",
  "kernelspec": {
   "display_name": "Python 3.9.13 ('base')",
   "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.9.13"
  },
  "toc-showcode": false,
  "vscode": {
   "interpreter": {
    "hash": "c7fbf5898c3415d81c1b4d47dfe00d3da846b26a9ac775122b4526fdc49a858f"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}