Colab_notebooks/StarDist_3D_ZeroCostDL4Mic.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IkSguVy8Xv83"
      },
      "source": [
        "# **StarDist (3D)**\n",
        "---\n",
        "\n",
        "<font size = 4>**StarDist 3D** is a deep-learning method that can be used to segment cell nuclei from 3D bioimages and was first published by [Weigert *et al.* in 2019 on arXiv](https://arxiv.org/abs/1908.03636), extending to 3D the 2D appraoch from [Schmidt *et al.* in 2018](https://arxiv.org/abs/1806.03535). It uses a shape representation based on star-convex polygons for nuclei in an image to predict the presence and the shape of these nuclei. This StarDist 3D network is based on an adapted ResNet network architecture.\n",
        "\n",
        "<font size = 4> **This particular notebook enables nuclei segmentation of 3D dataset. If you are interested in 2D dataset, you should use the StarDist 2D notebook instead.**\n",
        "\n",
        "---\n",
        "<font size = 4>*Disclaimer*:\n",
        "\n",
        "<font size = 4>This notebook is part of the Zero-Cost Deep-Learning to Enhance Microscopy project (https://github.com/HenriquesLab/DeepLearning_Collab/wiki). Jointly developed by the Jacquemet (link to https://cellmig.org/) and Henriques (https://henriqueslab.github.io/) laboratories.\n",
        "\n",
        "<font size = 4>This notebook is largely based on the paper:\n",
        "\n",
        "<font size = 4>**Cell Detection with Star-convex Polygons** from Schmidt *et al.*, International Conference on Medical Image Computing and Computer-Assisted Intervention (MICCAI), Granada, Spain, September 2018. (https://arxiv.org/abs/1806.03535)\n",
        "\n",
        "<font size = 4>and the 3D extension of the approach:\n",
        "\n",
        "<font size = 4>**Star-convex Polyhedra for 3D Object Detection and Segmentation in Microscopy** from Weigert *et al.* published on arXiv in 2019 (https://arxiv.org/abs/1908.03636)\n",
        "\n",
        "<font size = 4>**The Original code** is freely available in GitHub:\n",
        "https://github.com/mpicbg-csbd/stardist\n",
        "\n",
        "<font size = 4>**Please also cite this original paper when using or developing this notebook.**\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jWAz2i7RdxUV"
      },
      "source": [
        "# **How to use this notebook?**\n",
        "\n",
        "---\n",
        "\n",
        "<font size = 4>Video describing how to use our notebooks are available on youtube:\n",
        "  - [**Video 1**](https://www.youtube.com/watch?v=GzD2gamVNHI&feature=youtu.be): Full run through of the workflow to obtain the notebooks and the provided test datasets as well as a common use of the notebook\n",
        "  - [**Video 2**](https://www.youtube.com/watch?v=PUuQfP5SsqM&feature=youtu.be): Detailed description of the different sections of the notebook\n",
        "\n",
        "\n",
        "---\n",
        "###**Structure of a notebook**\n",
        "\n",
        "<font size = 4>The notebook contains two types of cell:  \n",
        "\n",
        "<font size = 4>**Text cells** provide information and can be modified by douple-clicking the cell. You are currently reading the text cell. You can create a new text by clicking `+ Text`.\n",
        "\n",
        "<font size = 4>**Code cells** contain code and the code can be modfied by selecting the cell. To execute the cell, move your cursor on the `[ ]`-mark on the left side of the cell (play button appears). Click to execute the cell. After execution is done the animation of play button stops. You can create a new coding cell by clicking `+ Code`.\n",
        "\n",
        "---\n",
        "###**Table of contents, Code snippets** and **Files**\n",
        "\n",
        "<font size = 4>On the top left side of the notebook you find three tabs which contain from top to bottom:\n",
        "\n",
        "<font size = 4>*Table of contents* = contains structure of the notebook. Click the content to move quickly between sections.\n",
        "\n",
        "<font size = 4>*Code snippets* = contain examples how to code certain tasks. You can ignore this when using this notebook.\n",
        "\n",
        "<font size = 4>*Files* = contain all available files. After mounting your google drive (see section 1.) you will find your files and folders here. \n",
        "\n",
        "<font size = 4>**Remember that all uploaded files are purged after changing the runtime.** All files saved in Google Drive will remain. You do not need to use the Mount Drive-button; your Google Drive is connected in section 1.2.\n",
        "\n",
        "<font size = 4>**Note:** The \"sample data\" in \"Files\" contains default files. Do not upload anything in here!\n",
        "\n",
        "---\n",
        "###**Making changes to the notebook**\n",
        "\n",
        "<font size = 4>**You can make a copy** of the notebook and save it to your Google Drive. To do this click file -> save a copy in drive.\n",
        "\n",
        "<font size = 4>To **edit a cell**, double click on the text. This will show you either the source code (in code cells) or the source text (in text cells).\n",
        "You can use the `#`-mark in code cells to comment out parts of the code. This allows you to keep the original code piece in the cell as a comment."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gKDLkLWUd-YX"
      },
      "source": [
        "#**0. Before getting started**\n",
        "---\n",
        "<font size = 4> For StarDist to train, **it needs to have access to a paired training dataset made of images of nuclei and their corresponding masks**. Information on how to generate a training dataset is available in our Wiki page: https://github.com/HenriquesLab/ZeroCostDL4Mic/wiki\n",
        "\n",
        "<font size = 4>**We strongly recommend that you generate extra paired images. These images can be used to assess the quality of your trained model (Quality control dataset)**. The quality control assessment can be done directly in this notebook.\n",
        "\n",
        "<font size = 4>The data structure is important. It is necessary that all the input data are in the same folder and that all the output data is in a separate folder. The provided training dataset is already split in two folders called \"Training - Images\" (Training_source) and \"Training - Masks\" (Training_target).\n",
        "\n",
        "<font size = 4>Additionally, the corresponding Training_source and Training_target files need to have **the same name**.\n",
        "\n",
        "<font size = 4>Please note that you currently can **only use .tif files!**\n",
        "\n",
        "<font size = 4>You can also provide a folder that contains the data that you wish to analyse with the trained network once all training has been performed.\n",
        "\n",
        "<font size = 4>Here's a common data structure that can work:\n",
        "*   Experiment A\n",
        "    - **Training dataset**\n",
        "      - Images of nuclei (Training_source)\n",
        "        - img_1.tif, img_2.tif, ...\n",
        "      - Masks (Training_target)\n",
        "        - img_1.tif, img_2.tif, ...\n",
        "    - **Quality control dataset**\n",
        "     - Images of nuclei\n",
        "        - img_1.tif, img_2.tif\n",
        "      - **Masks** \n",
        "        - img_1.tif, img_2.tif\n",
        "    - **Data to be predicted**\n",
        "    - **Results**\n",
        "\n",
        "---\n",
        "<font size = 4>**Important note**\n",
        "\n",
        "<font size = 4>- If you wish to **Train a network from scratch** using your own dataset (and we encourage everyone to do that), you will need to run **sections 1 - 4**, then use **section 5** to assess the quality of your model and **section 6** to run predictions using the model that you trained.\n",
        "\n",
        "<font size = 4>- If you wish to **Evaluate your model** using a model previously generated and saved on your Google Drive, you will only need to run **sections 1 and 2** to set up the notebook, then use **section 5** to assess the quality of your model.\n",
        "\n",
        "<font size = 4>- If you only wish to **run predictions** using a model previously generated and saved on your Google Drive, you will only need to run **sections 1 and 2** to set up the notebook, then use **section 6** to run the predictions on the desired model.\n",
        "---"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AdN8B91xZO0x"
      },
      "source": [
        "# **1. Install StarDist and dependencies**\n",
        "---\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "n16OlQru8MHl"
      },
      "source": [
        "## **1.1 Install StarDist**\n",
        "---\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "vvtlJacqSmOi"
      },
      "outputs": [],
      "source": [
        "#@markdown ##Install StarDist and dependencies\n",
        "\n",
        "\n",
        "!pip install -q tifffile # contains tools to operate tiff-files\n",
        "!pip install -q csbdeep  # contains tools for restoration of fluorescence microcopy images (Content-aware Image Restoration, CARE). It uses Keras and Tensorflow.\n",
        "#!pip install -e git://github.com/stardist/stardist.git@0.6.2#egg=stardist\n",
        "!pip install -q stardist==0.6.2\n",
        "!pip install -q gputools\n",
        "#!pip install edt\n",
        "!pip install -q wget\n",
        "!pip install -q fpdf2\n",
        "\n",
        "#Force session restart\n",
        "exit(0)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_IG7IqJtSuFY"
      },
      "source": [
        "## **1.2. Restart your runtime**\n",
        "---\n",
        "<font size = 4>\n",
        "\n",
        "\n",
        "**<font size = 4> Ignore the following message error message. Your Runtime has automatically restarted. This is normal.**\n",
        "\n",
        "<img width=\"40%\" alt =\"\" src=\"https://github.com/HenriquesLab/ZeroCostDL4Mic/raw/master/Wiki_files/session_crash.png\"><figcaption>  </figcaption>\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "40rqHN2PTJlu"
      },
      "source": [
        "## **1.3. Load key dependencies**\n",
        "---\n",
        "<font size = 4> "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "fq21zJVFNASx"
      },
      "outputs": [],
      "source": [
        "from __future__ import print_function, unicode_literals, absolute_import, division\n",
        "#@markdown ##Load key dependencies\n",
        "Notebook_version = '1.15.2'\n",
        "Network = 'StarDist 3D'\n",
        "\n",
        "from builtins import any as b_any\n",
        "\n",
        "def get_requirements_path():\n",
        "    # Store requirements file in 'contents' directory \n",
        "    current_dir = os.getcwd()\n",
        "    dir_count = current_dir.count('/') - 1\n",
        "    path = '../' * (dir_count) + 'requirements.txt'\n",
        "    return path\n",
        "\n",
        "def filter_files(file_list, filter_list):\n",
        "    filtered_list = []\n",
        "    for fname in file_list:\n",
        "        if b_any(fname.split('==')[0] in s for s in filter_list):\n",
        "            filtered_list.append(fname)\n",
        "    return filtered_list\n",
        "\n",
        "def build_requirements_file(before, after):\n",
        "    path = get_requirements_path()\n",
        "\n",
        "    # Exporting requirements.txt for local run\n",
        "    !pip freeze > $path\n",
        "\n",
        "    # Get minimum requirements file\n",
        "    df = pd.read_csv(path)\n",
        "    mod_list = [m.split('.')[0] for m in after if not m in before]\n",
        "    req_list_temp = df.values.tolist()\n",
        "    req_list = [x[0] for x in req_list_temp]\n",
        "\n",
        "    # Replace with package name and handle cases where import name is different to module name\n",
        "    mod_name_list = [['sklearn', 'scikit-learn'], ['skimage', 'scikit-image']]\n",
        "    mod_replace_list = [[x[1] for x in mod_name_list] if s in [x[0] for x in mod_name_list] else s for s in mod_list] \n",
        "    filtered_list = filter_files(req_list, mod_replace_list)\n",
        "\n",
        "    file=open(path,'w')\n",
        "    for item in filtered_list:\n",
        "        file.writelines(item)\n",
        "\n",
        "    file.close()\n",
        "\n",
        "import sys\n",
        "before = [str(m) for m in sys.modules]\n",
        "\n",
        "# %tensorflow_version 1.x\n",
        "import tensorflow\n",
        "print(tensorflow.__version__)\n",
        "print(\"Tensorflow enabled.\")\n",
        "\n",
        "# ------- Variable specific to Stardist -------\n",
        "from stardist import fill_label_holes, random_label_cmap, calculate_extents, gputools_available\n",
        "from stardist.models import Config3D, StarDist3D, StarDistData3D\n",
        "from stardist import relabel_image_stardist3D, Rays_GoldenSpiral, calculate_extents\n",
        "from stardist.matching import matching_dataset\n",
        "from csbdeep.utils import Path, normalize, download_and_extract_zip_file, plot_history # for loss plot\n",
        "from csbdeep.io import save_tiff_imagej_compatible\n",
        "import numpy as np\n",
        "np.random.seed(42)\n",
        "lbl_cmap = random_label_cmap()\n",
        "import cv2\n",
        "%matplotlib inline\n",
        "%config InlineBackend.figure_format = 'retina'\n",
        "\n",
        "# ------- Common variable to all ZeroCostDL4Mic notebooks -------\n",
        "import numpy as np\n",
        "from matplotlib import pyplot as plt\n",
        "import urllib\n",
        "import os, random\n",
        "import shutil \n",
        "import zipfile\n",
        "from tifffile import imread, imsave\n",
        "from csbdeep.io import save_tiff_imagej_compatible\n",
        "import time\n",
        "import sys\n",
        "import wget\n",
        "from pathlib import Path\n",
        "import pandas as pd\n",
        "import csv\n",
        "from glob import glob\n",
        "from scipy import signal\n",
        "from scipy import ndimage\n",
        "from skimage import io\n",
        "from sklearn.linear_model import LinearRegression\n",
        "from skimage.util import img_as_uint\n",
        "import matplotlib as mpl\n",
        "from skimage.metrics import structural_similarity\n",
        "from skimage.metrics import peak_signal_noise_ratio as psnr\n",
        "from astropy.visualization import simple_norm\n",
        "from skimage import img_as_float32\n",
        "from skimage.util import img_as_ubyte\n",
        "from tqdm import tqdm \n",
        "from fpdf import FPDF, HTMLMixin\n",
        "from datetime import datetime\n",
        "from pip._internal.operations.freeze import freeze\n",
        "import subprocess\n",
        "from astropy.visualization import simple_norm\n",
        "\n",
        "\n",
        "# For sliders and dropdown menu and progress bar\n",
        "from ipywidgets import interact\n",
        "import ipywidgets as widgets\n",
        "\n",
        "# Colors for the warning messages\n",
        "class bcolors:\n",
        "  WARNING = '\\033[31m'\n",
        "\n",
        "W  = '\\033[0m'  # white (normal)\n",
        "R  = '\\033[31m' # red\n",
        "\n",
        "#Disable some of the tensorflow warnings\n",
        "import warnings\n",
        "warnings.filterwarnings(\"ignore\")\n",
        "\n",
        "print(\"Libraries installed\")\n",
        "\n",
        "\n",
        "\n",
        "# Check if this is the latest version of the notebook\n",
        "All_notebook_versions = pd.read_csv(\"https://raw.githubusercontent.com/HenriquesLab/ZeroCostDL4Mic/master/Colab_notebooks/Latest_Notebook_versions.csv\", dtype=str)\n",
        "print('Notebook version: '+Notebook_version)\n",
        "Latest_Notebook_version = All_notebook_versions[All_notebook_versions[\"Notebook\"] == Network]['Version'].iloc[0]\n",
        "print('Latest notebook version: '+Latest_Notebook_version)\n",
        "if Notebook_version == Latest_Notebook_version:\n",
        "  print(\"This notebook is up-to-date.\")\n",
        "else:\n",
        "  print(bcolors.WARNING +\"A new version of this notebook has been released. We recommend that you download it at https://github.com/HenriquesLab/ZeroCostDL4Mic/wiki\")\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "def pdf_export(trained = False, augmentation = False, pretrained_model = False):\n",
        "  class MyFPDF(FPDF, HTMLMixin):\n",
        "    pass\n",
        "\n",
        "  pdf = MyFPDF()\n",
        "  pdf.add_page()\n",
        "  pdf.set_right_margin(-1)\n",
        "  pdf.set_font(\"Arial\", size = 11, style='B') \n",
        "\n",
        "  \n",
        "  day = datetime.now()\n",
        "  datetime_str = str(day)[0:10]\n",
        "\n",
        "  Header = 'Training report for '+Network+' model ('+model_name+')\\nDate: '+datetime_str\n",
        "  pdf.multi_cell(180, 5, txt = Header, align = 'L') \n",
        "  pdf.ln(1)\n",
        "    \n",
        "  # add another cell \n",
        "  if trained:\n",
        "    training_time = \"Training time: \"+str(hour)+ \"hour(s) \"+str(mins)+\"min(s) \"+str(round(sec))+\"sec(s)\"\n",
        "    pdf.cell(190, 5, txt = training_time, ln = 1, align='L')\n",
        "  pdf.ln(1)\n",
        "\n",
        "  Header_2 = 'Information for your materials and methods:'\n",
        "  pdf.cell(190, 5, txt=Header_2, ln=1, align='L')\n",
        "\n",
        "  all_packages = ''\n",
        "  for requirement in freeze(local_only=True):\n",
        "    all_packages = all_packages+requirement+', '\n",
        "  #print(all_packages)\n",
        "\n",
        "  #Main Packages\n",
        "  main_packages = ''\n",
        "  version_numbers = []\n",
        "  for name in ['tensorflow','numpy','keras','csbdeep']:\n",
        "    find_name=all_packages.find(name)\n",
        "    main_packages = main_packages+all_packages[find_name:all_packages.find(',',find_name)]+', '\n",
        "    #Version numbers only here:\n",
        "    version_numbers.append(all_packages[find_name+len(name)+2:all_packages.find(',',find_name)])\n",
        "\n",
        "  cuda_version = subprocess.run('nvcc --version',stdout=subprocess.PIPE, shell=True)\n",
        "  cuda_version = cuda_version.stdout.decode('utf-8')\n",
        "  cuda_version = cuda_version[cuda_version.find(', V')+3:-1]\n",
        "  gpu_name = subprocess.run('nvidia-smi',stdout=subprocess.PIPE, shell=True)\n",
        "  gpu_name = gpu_name.stdout.decode('utf-8')\n",
        "  gpu_name = gpu_name[gpu_name.find('Tesla'):gpu_name.find('Tesla')+10]\n",
        "  #print(cuda_version[cuda_version.find(', V')+3:-1])\n",
        "  #print(gpu_name)\n",
        "\n",
        "  shape = io.imread(Training_source+'/'+os.listdir(Training_source)[1]).shape\n",
        "  dataset_size = len(os.listdir(Training_source))\n",
        "\n",
        "  text = 'The '+Network+' model was trained from scratch for '+str(number_of_epochs)+' epochs on '+str(int(len(X)))+' paired image patches (image dimensions: '+str(shape)+', patch size: ('+str(patch_height)+','+str(patch_size)+','+str(patch_size)+')) with a batch size of '+str(batch_size)+' and a '+conf.train_dist_loss+' loss function, using the '+Network+' ZeroCostDL4Mic notebook (v '+Notebook_version[0]+') (von Chamier & Laine et al., 2020). Key python packages used include tensorflow (v '+version_numbers[0]+'), keras (v '+version_numbers[2]+'), csbdeep (v '+version_numbers[3]+'), numpy (v '+version_numbers[1]+'), cuda (v '+cuda_version+'). The training was accelerated using a '+gpu_name+'GPU.'\n",
        "\n",
        "  #text = 'The '+Network+' model ('+model_name+') was trained using '+str(dataset_size)+' paired images (image dimensions: '+str(shape)+') using the '+Network+' ZeroCostDL4Mic notebook (v '+Notebook_version[0]+') (von Chamier & Laine et al., 2020). Key python packages used include tensorflow (v '+version_numbers[0]+'), Keras (v '+version_numbers[2]+'), csbdeep (v '+version_numbers[3]+'), numpy (v '+version_numbers[1]+'), cuda (v '+cuda_version+'). The GPU used was a '+gpu_name+'.'\n",
        "\n",
        "  if pretrained_model:\n",
        "    text = 'The '+Network+' model was trained for '+str(number_of_epochs)+' epochs on '+str(int(len(X)))+' paired image patches (image dimensions: '+str(shape)+', patch size: ('+str(patch_height)+','+str(patch_size)+','+str(patch_size)+')) with a batch size of '+str(batch_size)+' and a '+conf.train_dist_loss+' loss function, using the '+Network+' ZeroCostDL4Mic notebook (v '+Notebook_version[0]+') (von Chamier & Laine et al., 2020). The model was retrained from a pretrained model. Key python packages used include tensorflow (v '+version_numbers[0]+'), keras (v '+version_numbers[2]+'), csbdeep (v '+version_numbers[3]+'), numpy (v '+version_numbers[1]+'), cuda (v '+cuda_version+'). The training was accelerated using a '+gpu_name+'GPU.'\n",
        "\n",
        "  pdf.set_font('')\n",
        "  pdf.set_font_size(10.)\n",
        "  pdf.multi_cell(190, 5, txt = text, align='L')\n",
        "  pdf.ln(1)\n",
        "  pdf.set_font('')\n",
        "  pdf.set_font('Arial', size = 10, style = 'B')\n",
        "  pdf.cell(28, 5, txt='Augmentation: ', ln=0)\n",
        "  pdf.set_font('')\n",
        "  if augmentation:\n",
        "    aug_text = 'The dataset was augmented.'\n",
        "  else:\n",
        "    aug_text = 'No augmentation was used for training.'\n",
        "  pdf.multi_cell(190, 5, txt=aug_text, align='L')\n",
        "  pdf.ln(1)\n",
        "  pdf.set_font('Arial', size = 11, style = 'B')\n",
        "  pdf.cell(180, 5, txt = 'Parameters', align='L', ln=1)\n",
        "  pdf.set_font('')\n",
        "  pdf.set_font_size(10.)\n",
        "  if Use_Default_Advanced_Parameters:\n",
        "    pdf.cell(200, 5, txt='Default Advanced Parameters were enabled')\n",
        "  pdf.cell(200, 5, txt='The following parameters were used for training:')\n",
        "  pdf.ln(1)\n",
        "  html = \"\"\" \n",
        "  <table width=40% style=\"margin-left:0px;\">\n",
        "    <tr>\n",
        "      <th width = 50% align=\"left\">Parameter</th>\n",
        "      <th width = 50% align=\"left\">Value</th>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <td width = 50%>number_of_epochs</td>\n",
        "      <td width = 50%>{0}</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <td width = 50%>patch_size</td>\n",
        "      <td width = 50%>{1}</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <td width = 50%>batch_size</td>\n",
        "      <td width = 50%>{2}</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <td width = 50%>number_of_steps</td>\n",
        "      <td width = 50%>{3}</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <td width = 50%>percentage_validation</td>\n",
        "      <td width = 50%>{4}</td>\n",
        "    </tr>\n",
        "      <tr>\n",
        "      <td width = 50%>n_rays</td>\n",
        "      <td width = 50%>{5}</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <td width = 50%>initial_learning_rate</td>\n",
        "      <td width = 50%>{6}</td>\n",
        "    </tr>\n",
        "  </table>\n",
        "  \"\"\".format(number_of_epochs,str(patch_height)+'x'+str(patch_size)+'x'+str(patch_size),batch_size,number_of_steps,percentage_validation,n_rays,initial_learning_rate)\n",
        "  pdf.write_html(html)\n",
        "\n",
        "  #pdf.multi_cell(190, 5, txt = text_2, align='L')\n",
        "  pdf.set_font(\"Arial\", size = 11, style='B')\n",
        "  pdf.ln(1)\n",
        "  pdf.cell(190, 5, txt = 'Training Dataset', align='L', ln=1)\n",
        "  pdf.set_font('')\n",
        "  pdf.set_font('Arial', size = 10, style = 'B')\n",
        "  pdf.cell(30, 5, txt= 'Training_source:', align = 'L', ln=0)\n",
        "  pdf.set_font('')\n",
        "  pdf.multi_cell(170, 5, txt = Training_source, align = 'L')\n",
        "  pdf.ln(1)\n",
        "  pdf.set_font('')\n",
        "  pdf.set_font('Arial', size = 10, style = 'B')\n",
        "  pdf.cell(28, 5, txt= 'Training_target:', align = 'L', ln=0)\n",
        "  pdf.set_font('')\n",
        "  pdf.multi_cell(170, 5, txt = Training_target, align = 'L')\n",
        "  pdf.ln(1)\n",
        "  pdf.set_font('')\n",
        "  pdf.set_font('Arial', size = 10, style = 'B')\n",
        "  pdf.cell(21, 5, txt= 'Model Path:', align = 'L', ln=0)\n",
        "  pdf.set_font('')\n",
        "  pdf.multi_cell(170, 5, txt = model_path+'/'+model_name, align = 'L')\n",
        "  pdf.ln(1)\n",
        "  pdf.cell(60, 5, txt = 'Example Training pair', ln=1)\n",
        "  pdf.ln(1)\n",
        "  exp_size = io.imread('/content/TrainingDataExample_StarDist3D.png').shape\n",
        "  pdf.image('/content/TrainingDataExample_StarDist3D.png', x = 11, y = None, w = round(exp_size[1]/8), h = round(exp_size[0]/8))\n",
        "  pdf.ln(1)\n",
        "  ref_1 = 'References:\\n - ZeroCostDL4Mic: von Chamier, Lucas & Laine, Romain, et al. \"Democratising deep learning for microscopy with ZeroCostDL4Mic.\" Nature Communications (2021).'\n",
        "  pdf.multi_cell(190, 5, txt = ref_1, align='L')\n",
        "  pdf.ln(1)\n",
        "  ref_2 = '- StarDist 3D: Schmidt, Uwe, et al. \"Cell detection with star-convex polygons.\" International Conference on Medical Image Computing and Computer-Assisted Intervention. Springer, Cham, 2018.'\n",
        "  pdf.multi_cell(190, 5, txt = ref_2, align='L')\n",
        "  pdf.ln(1)\n",
        "  ref_3 = '- StarDist 3D: Weigert, Martin, et al. \"Star-convex polyhedra for 3d object detection and segmentation in microscopy.\" The IEEE Winter Conference on Applications of Computer Vision. 2020.'\n",
        "  pdf.multi_cell(190, 5, txt = ref_3, align='L')\n",
        "  # if Use_Data_augmentation:\n",
        "  #   ref_4 = '- Augmentor: Bloice, Marcus D., Christof Stocker, and Andreas Holzinger. \"Augmentor: an image augmentation library for machine learning.\" arXiv preprint arXiv:1708.04680 (2017).'\n",
        "  #   pdf.multi_cell(190, 5, txt = ref_4, align='L')\n",
        "  pdf.ln(3)\n",
        "  reminder = 'Important:\\nRemember to perform the quality control step on all newly trained models\\nPlease consider depositing your training dataset on Zenodo'\n",
        "  pdf.set_font('Arial', size = 11, style='B')\n",
        "  pdf.multi_cell(190, 5, txt=reminder, align='C')\n",
        "  pdf.ln(1)\n",
        "\n",
        "  pdf.output(model_path+'/'+model_name+'/'+model_name+\"_training_report.pdf\")\n",
        "\n",
        "\n",
        "def qc_pdf_export():\n",
        "  class MyFPDF(FPDF, HTMLMixin):\n",
        "    pass\n",
        "\n",
        "  pdf = MyFPDF()\n",
        "  pdf.add_page()\n",
        "  pdf.set_right_margin(-1)\n",
        "  pdf.set_font(\"Arial\", size = 11, style='B') \n",
        "\n",
        "  Network = 'Stardist 3D'\n",
        "\n",
        "  day = datetime.now()\n",
        "  datetime_str = str(day)[0:10]\n",
        "\n",
        "  Header = 'Quality Control report for '+Network+' model ('+QC_model_name+')\\nDate: '+datetime_str\n",
        "  pdf.multi_cell(180, 5, txt = Header, align = 'L') \n",
        "  pdf.ln(1)\n",
        "\n",
        "  all_packages = ''\n",
        "  for requirement in freeze(local_only=True):\n",
        "    all_packages = all_packages+requirement+', '\n",
        "\n",
        "  pdf.set_font('')\n",
        "  pdf.set_font('Arial', size = 11, style = 'B')\n",
        "  pdf.ln(2)\n",
        "  pdf.cell(190, 5, txt = 'Development of Training Losses', ln=1, align='L')\n",
        "  pdf.ln(1)\n",
        "  exp_size = io.imread(full_QC_model_path+'/Quality Control/lossCurvePlots.png').shape\n",
        "  if os.path.exists(full_QC_model_path+'/Quality Control/lossCurvePlots.png'):\n",
        "    pdf.image(full_QC_model_path+'/Quality Control/lossCurvePlots.png', x = 11, y = None, w = round(exp_size[1]/8), h = round(exp_size[0]/8))\n",
        "  else:\n",
        "    pdf.set_font('')\n",
        "    pdf.set_font('Arial', size=10)\n",
        "    pdf.multi_cell(190, 5, txt='If you would like to see the evolution of the loss function during training please play the first cell of the QC section in the notebook.')\n",
        "  pdf.ln(2)\n",
        "  pdf.set_font('')\n",
        "  pdf.set_font('Arial', size = 10, style = 'B')\n",
        "  pdf.ln(3)\n",
        "  pdf.cell(80, 5, txt = 'Example Quality Control Visualisation', ln=1)\n",
        "  pdf.ln(1)\n",
        "  exp_size = io.imread(full_QC_model_path+'/Quality Control/QC_example_data.png').shape\n",
        "  pdf.image(full_QC_model_path+'/Quality Control/QC_example_data.png', x = 16, y = None, w = round(exp_size[1]/10), h = round(exp_size[0]/10))\n",
        "  pdf.ln(1)\n",
        "  pdf.set_font('')\n",
        "  pdf.set_font('Arial', size = 11, style = 'B')\n",
        "  pdf.ln(1)\n",
        "  pdf.cell(180, 5, txt = 'Quality Control Metrics', align='L', ln=1)\n",
        "  pdf.set_font('')\n",
        "  pdf.set_font_size(10.)\n",
        "\n",
        "  pdf.ln(1)\n",
        "  html = \"\"\"\n",
        "  <body>\n",
        "  <font size=\"10\" face=\"Courier\" >\n",
        "  <table width=50% style=\"margin-left:0px;\">\"\"\"\n",
        "  with open(full_QC_model_path+'/Quality Control/Quality_Control for '+QC_model_name+'.csv', 'r') as csvfile:\n",
        "    metrics = csv.reader(csvfile)\n",
        "    header = next(metrics)\n",
        "    image = header[0]\n",
        "    PvGT_IoU = header[1]\n",
        "    header = \"\"\"\n",
        "    <tr>\n",
        "    <th width = 50% align=\"center\">{0}</th>\n",
        "    <th width = 50% align=\"center\">{1}</th>\n",
        "    </tr>\"\"\".format(image,PvGT_IoU)\n",
        "    html = html+header\n",
        "    for row in metrics:\n",
        "      image = row[0]\n",
        "      PvGT_IoU = row[1]\n",
        "      cells = \"\"\"\n",
        "        <tr>\n",
        "          <td width = 50% align=\"center\">{0}</td>\n",
        "          <td width = 50% align=\"center\">{1}</td>\n",
        "        </tr>\"\"\".format(image,str(round(float(PvGT_IoU),3)))\n",
        "      html = html+cells\n",
        "    html = html+\"\"\"</body></table>\"\"\"\n",
        "    \n",
        "  pdf.write_html(html)\n",
        "\n",
        "  pdf.ln(1)\n",
        "  pdf.set_font('')\n",
        "  pdf.set_font_size(10.)\n",
        "  ref_1 = 'References:\\n - ZeroCostDL4Mic: von Chamier, Lucas & Laine, Romain, et al. \"Democratising deep learning for microscopy with ZeroCostDL4Mic.\" Nature Communications (2021).'\n",
        "  pdf.multi_cell(190, 5, txt = ref_1, align='L')\n",
        "  pdf.ln(1)\n",
        "  ref_2 = ' - Weigert, Martin, et al. \"Star-convex polyhedra for 3d object detection and segmentation in microscopy.\" The IEEE Winter Conference on Applications of Computer Vision. 2020.'\n",
        "  pdf.multi_cell(190, 5, txt = ref_2, align='L')\n",
        "  pdf.ln(1)\n",
        "\n",
        "  pdf.ln(3)\n",
        "  reminder = 'To find the parameters and other information about how this model was trained, go to the training_report.pdf of this model which should be in the folder of the same name.'\n",
        "\n",
        "  pdf.set_font('Arial', size = 11, style='B')\n",
        "  pdf.multi_cell(190, 5, txt=reminder, align='C')\n",
        "  pdf.ln(1)\n",
        "\n",
        "  pdf.output(full_QC_model_path+'/Quality Control/'+QC_model_name+'_QC_report.pdf')\n",
        "\n",
        "\n",
        "\n",
        "def random_fliprot(img, mask, axis=None): \n",
        "    if axis is None:\n",
        "        axis = tuple(range(mask.ndim))\n",
        "    axis = tuple(axis)\n",
        "            \n",
        "    assert img.ndim>=mask.ndim\n",
        "    perm = tuple(np.random.permutation(axis))\n",
        "    transpose_axis = np.arange(mask.ndim)\n",
        "    for a, p in zip(axis, perm):\n",
        "        transpose_axis[a] = p\n",
        "    transpose_axis = tuple(transpose_axis)\n",
        "    img = img.transpose(transpose_axis + tuple(range(mask.ndim, img.ndim))) \n",
        "    mask = mask.transpose(transpose_axis) \n",
        "    for ax in axis: \n",
        "        if np.random.rand() > 0.5:\n",
        "            img = np.flip(img, axis=ax)\n",
        "            mask = np.flip(mask, axis=ax)\n",
        "    return img, mask \n",
        "\n",
        "def random_intensity_change(img):\n",
        "    img = img*np.random.uniform(0.6,2) + np.random.uniform(-0.2,0.2)\n",
        "    return img\n",
        "\n",
        "def augmenter(x, y):\n",
        "    \"\"\"Augmentation of a single input/label image pair.\n",
        "    x is an input image\n",
        "    y is the corresponding ground-truth label image\n",
        "    \"\"\"\n",
        "    # Note that we only use fliprots along axis=(1,2), i.e. the yx axis \n",
        "    # as 3D microscopy acquisitions are usually not axially symmetric\n",
        "    x, y = random_fliprot(x, y, axis=(1,2))\n",
        "    x = random_intensity_change(x)\n",
        "    return x, y\n",
        "\n",
        "# Build requirements file for local run\n",
        "after = [str(m) for m in sys.modules]\n",
        "build_requirements_file(before, after)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "n4yWFoJNnoin"
      },
      "source": [
        "# **2. Initialise the Colab session**\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "---\n",
        "\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DMNHVZfHmbKb"
      },
      "source": [
        "\n",
        "## **2.1. Check for GPU access**\n",
        "---\n",
        "\n",
        "By default, the session should be using Python 3 and GPU acceleration, but it is possible to ensure that these are set properly by doing the following:\n",
        "\n",
        "<font size = 4>Go to **Runtime -> Change the Runtime type**\n",
        "\n",
        "<font size = 4>**Runtime type: Python 3** *(Python 3 is programming language in which this program is written)*\n",
        "\n",
        "<font size = 4>**Accelerator: GPU** *(Graphics processing unit)*\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "zCvebubeSaGY"
      },
      "outputs": [],
      "source": [
        "#@markdown ##Run this cell to check if you have GPU access\n",
        "# %tensorflow_version 1.x\n",
        "\n",
        "\n",
        "import tensorflow as tf\n",
        "if tf.test.gpu_device_name()=='':\n",
        "  print('You do not have GPU access.') \n",
        "  print('Did you change your runtime ?') \n",
        "  print('If the runtime setting is correct then Google did not allocate a GPU for your session')\n",
        "  print('Expect slow performance. To access GPU try reconnecting later')\n",
        "\n",
        "else:\n",
        "  print('You have GPU access')\n",
        "  !nvidia-smi"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sNIVx8_CLolt"
      },
      "source": [
        "## **2.2. Mount your Google Drive**\n",
        "---\n",
        "<font size = 4> To use this notebook on the data present in your Google Drive, you need to mount your Google Drive to this notebook.\n",
        "\n",
        "<font size = 4> Play the cell below to mount your Google Drive and follow the link. In the new browser window, select your drive and select 'Allow', copy the code, paste into the cell and press enter. This will give Colab access to the data on the drive. \n",
        "\n",
        "<font size = 4> Once this is done, your data are available in the **Files** tab on the top left of notebook."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "01Djr8v-5pPk"
      },
      "outputs": [],
      "source": [
        "#@markdown ##Play the cell to connect your Google Drive to Colab\n",
        "\n",
        "#@markdown * Click on the URL. \n",
        "\n",
        "#@markdown * Sign in your Google Account. \n",
        "\n",
        "#@markdown * Copy the authorization code. \n",
        "\n",
        "#@markdown * Enter the authorization code. \n",
        "\n",
        "#@markdown * Click on \"Files\" site on the right. Refresh the site. Your Google Drive folder should now be available here as \"drive\". \n",
        "\n",
        "# mount user's Google Drive to Google Colab.\n",
        "from google.colab import drive\n",
        "drive.mount('/content/gdrive')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7ncSG54e9LU4"
      },
      "source": [
        "**<font size = 4> If you cannot see your files, reactivate your session by connecting to your hosted runtime.** \n",
        "\n",
        "\n",
        "<img width=\"40%\" alt =\"Example of image detection with retinanet.\" src=\"https://github.com/HenriquesLab/ZeroCostDL4Mic/raw/master/Wiki_files/connect_to_hosted.png\"><figcaption> Connect to a hosted runtime. </figcaption>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HLYcZR9gMv42"
      },
      "source": [
        "# **3. Select your parameters and paths**\n",
        "---"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FQ_QxtSWQ7CL"
      },
      "source": [
        "## **3.1. Setting main training parameters**\n",
        "---\n",
        "<font size = 4> "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AuESFimvMv43"
      },
      "source": [
        "<font size = 5> **Paths for training, predictions and results**\n",
        "\n",
        "<font size = 4>**`Training_source:`, `Training_target`:** These are the paths to your folders containing the Training_source (images of nuclei) and Training_target (masks) training data respecively. To find the paths of the folders containing the respective datasets, go to your Files on the left of the notebook, navigate to the folder containing your files and copy the path by right-clicking on the folder, **Copy path** and pasting it into the right box below.\n",
        "\n",
        "<font size = 4>**`model_name`:** Use only my_model -style, not my-model (Use \"_\" not \"-\"). Do not use spaces in the name. Avoid using the name of an existing model (saved in the same folder) as it will be overwritten.\n",
        "\n",
        "<font size = 4>**`model_path`**: Enter the path where your model will be saved once trained (for instance your result folder).\n",
        "\n",
        "<font size = 5>**Training parameters**\n",
        "\n",
        "<font size = 4>**`number_of_epochs`:** Input how many epochs (rounds) the network will be trained. Preliminary results can already be observed after a 200 epochs, but a full training should run for more. Evaluate the performance after training (see 5.). **Default value: 200**\n",
        "\n",
        "<font size = 5>**Advanced parameters - experienced users only**\n",
        "\n",
        "<font size =4>**`batch_size:`** This parameter defines the number of patches seen in each training step. Reducing or increasing the **batch size** may slow or speed up your training, respectively, and can influence network performance. **Default value: 1** \n",
        "\n",
        "<font size = 4>**`number_of_steps`:** Define the number of training steps by epoch. By default (or if set to 0) this parameter is calculated so that each image / patch is seen at least once per epoch. **Default value: Number of patch / batch_size**. This value is multiplied by 6 when augmentation is enabled.\n",
        "\n",
        "<font size = 4>**`patch_size`:** and **`patch_height`:** Input the size of the patches use to train StarDist 3D (length of a side). The value should be smaller or equal to the dimensions of the image. Make patch size and patch_height as large as possible and divisible by 8 and 4, respectively. **Default value: dimension of the training images**\n",
        "\n",
        "<font size = 4>**`percentage_validation`:**  Input the percentage of your training dataset you want to use to validate the network during the training. **Default value: 10** \n",
        "\n",
        "<font size = 4>**`n_rays`:** Set number of rays (corners) used for StarDist (for instance a cube has 8 corners). **Default value: 96** \n",
        "\n",
        "<font size = 4>**`initial_learning_rate`:** Input the initial value to be used as learning rate. **Default value: 0.0003**\n",
        "\n",
        "<font size = 4>**If you get an Out of memory (OOM) error during the training,  manually decrease the patch_size and  patch_height values until the OOM error disappear.**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "ewpNJ_I0Mv47"
      },
      "outputs": [],
      "source": [
        "\n",
        "\n",
        "#@markdown ###Path to training images: \n",
        "Training_source = \"\" #@param {type:\"string\"}\n",
        "training_images = Training_source\n",
        "\n",
        "\n",
        "Training_target = \"\" #@param {type:\"string\"}\n",
        "mask_images = Training_target \n",
        "\n",
        "\n",
        "#@markdown ###Name of the model and path to model folder:\n",
        "model_name = \"\" #@param {type:\"string\"}\n",
        "\n",
        "model_path = \"\" #@param {type:\"string\"}\n",
        "trained_model = model_path \n",
        "\n",
        "#@markdown ### Other parameters for training:\n",
        "number_of_epochs =  100#@param {type:\"number\"}\n",
        "\n",
        "#@markdown ###Advanced Parameters\n",
        "Use_Default_Advanced_Parameters = True #@param {type:\"boolean\"}\n",
        "\n",
        "#@markdown ###If not, please input:\n",
        "\n",
        "#GPU_limit = 90 #@param {type:\"number\"}\n",
        "batch_size =  2#@param {type:\"number\"}\n",
        "number_of_steps = 0#@param {type:\"number\"}\n",
        "patch_size =  96#@param {type:\"number\"} # pixels in\n",
        "patch_height =  48#@param {type:\"number\"}\n",
        "percentage_validation =   10#@param {type:\"number\"}\n",
        "n_rays = 96 #@param {type:\"number\"}\n",
        "initial_learning_rate = 0.0003 #@param {type:\"number\"}\n",
        "\n",
        "if (Use_Default_Advanced_Parameters): \n",
        "  print(\"Default advanced parameters enabled\")\n",
        "  batch_size = 2 # default from original author's notebook\n",
        "  n_rays = 96\n",
        "  percentage_validation = 10\n",
        "  initial_learning_rate = 0.0003\n",
        "\n",
        "  patch_size =  96 # default from original author's notebook\n",
        "  patch_height =  48 # default from original author's notebook\n",
        "\n",
        "\n",
        "percentage = percentage_validation/100\n",
        "\n",
        "#here we check that no model with the same name already exist, if so print a warning\n",
        "\n",
        "if os.path.exists(model_path+'/'+model_name):\n",
        "  print(bcolors.WARNING +\"!! WARNING: \"+model_name+\" already exists and will be deleted in the following cell !!\")\n",
        "  print(bcolors.WARNING +\"To continue training \"+model_name+\", choose a new model_name here, and load \"+model_name+\" in section 3.3\"+W)\n",
        "   \n",
        "\n",
        "random_choice=random.choice(os.listdir(Training_source))\n",
        "x = imread(Training_source+\"/\"+random_choice)\n",
        "\n",
        "# Here we check that the input images are stacks\n",
        "if len(x.shape) == 3:\n",
        "  print(\"Image dimensions (z,y,x)\",x.shape)\n",
        "\n",
        "if not len(x.shape) == 3:\n",
        "  print(bcolors.WARNING +\"Your images appear to have the wrong dimensions. Image dimension\",x.shape)\n",
        "\n",
        "\n",
        "#Find image Z dimension and select the mid-plane\n",
        "Image_Z = x.shape[0]\n",
        "mid_plane = int(Image_Z / 2)+1\n",
        "\n",
        "\n",
        "#Find image XY dimension\n",
        "Image_Y = x.shape[1]\n",
        "Image_X = x.shape[2]\n",
        "\n",
        "# If default parameters, patch size is the same as image size\n",
        "# if (Use_Default_Advanced_Parameters): \n",
        "#   patch_size = min(Image_Y, Image_X) \n",
        "#   patch_height = Image_Z\n",
        "\n",
        "#Hyperparameters failsafes\n",
        "\n",
        "# Here we check that patch_size is smaller than the smallest xy dimension of the image \n",
        "\n",
        "if patch_size > min(Image_Y, Image_X):\n",
        "  patch_size = min(Image_Y, Image_X)\n",
        "  print (bcolors.WARNING + \" Your chosen patch_size is bigger than the xy dimension of your image; therefore the patch_size chosen is now:\",patch_size)\n",
        "\n",
        "# Here we check that patch_size is divisible by 8\n",
        "if not patch_size % 8 == 0:\n",
        "    patch_size = ((int(patch_size / 8)-1) * 8)\n",
        "    print (bcolors.WARNING + \" Your chosen patch_size is not divisible by 8; therefore the patch_size chosen is now:\",patch_size)\n",
        "\n",
        "# Here we check that patch_height is smaller than the z dimension of the image \n",
        "\n",
        "if patch_height > Image_Z :\n",
        "  patch_height = Image_Z\n",
        "  print (bcolors.WARNING + \" Your chosen patch_height is bigger than the z dimension of your image; therefore the patch_size chosen is now:\",patch_height)\n",
        "\n",
        "# Here we check that patch_height is divisible by 4\n",
        "if not patch_height % 4 == 0:\n",
        "    patch_height = ((int(patch_height / 4)-1) * 4)\n",
        "    if patch_height == 0:\n",
        "      patch_height = 4\n",
        "    print (bcolors.WARNING + \" Your chosen patch_height is not divisible by 4; therefore the patch_size chosen is now:\",patch_height)\n",
        "\n",
        "# Here we disable pre-trained model by default (in case the next cell is not ran)\n",
        "Use_pretrained_model = False\n",
        "\n",
        "# Here we disable data augmentation by default (in case the cell is not ran)\n",
        "\n",
        "Use_Data_augmentation = False\n",
        "\n",
        "print(\"Parameters initiated.\")\n",
        "\n",
        "\n",
        "os.chdir(Training_target)\n",
        "y = imread(Training_target+\"/\"+random_choice)\n",
        "\n",
        "#Here we use a simple normalisation strategy to visualise the image\n",
        "from astropy.visualization import simple_norm\n",
        "norm = simple_norm(x, percent = 99)\n",
        "\n",
        "mid_plane = int(Image_Z / 2)+1\n",
        "\n",
        "f=plt.figure(figsize=(16,8))\n",
        "plt.subplot(1,2,1)\n",
        "plt.imshow(x[mid_plane], interpolation='nearest', norm=norm, cmap='magma')\n",
        "plt.axis('off')\n",
        "plt.title('Training source (single Z plane)');\n",
        "plt.subplot(1,2,2)\n",
        "plt.imshow(y[mid_plane], interpolation='nearest', cmap=lbl_cmap)\n",
        "plt.axis('off')\n",
        "plt.title('Training target (single Z plane)');\n",
        "plt.savefig('/content/TrainingDataExample_StarDist3D.png',bbox_inches='tight',pad_inches=0)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xyQZKby8yFME"
      },
      "source": [
        "## **3.2. Data augmentation**\n",
        "---\n",
        "<font size = 4>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "w_jCy7xOx2g3"
      },
      "source": [
        "<font size = 4>Data augmentation can improve training progress by amplifying differences in the dataset. This can be useful if the available dataset is small since, in this case, it is possible that a network could quickly learn every example in the dataset (overfitting), without augmentation. Augmentation is not necessary for training and if your training dataset is large you should disable it.\n",
        "\n",
        "<font size = 4><font size = 4> **However, data augmentation is not a magic solution and may also introduce issues. Therefore, we recommend that you train your network with and without augmentation, and use the QC section to validate that it improves overall performances.**  \n",
        "\n",
        "<font size = 4>Data augmentation is performed here by flipping, rotating and modifying the intensity of the images.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "2ArWgdghVlQw"
      },
      "outputs": [],
      "source": [
        "#Data augmentation\n",
        "\n",
        "Use_Data_augmentation = False #@param {type:\"boolean\"}\n",
        "\n",
        "if Use_Data_augmentation:\n",
        "  augmenter = augmenter\n",
        "  print(\"Data augmentation enabled. Let's flip!\")\n",
        "else:\n",
        "  augmenter = None\n",
        "  print(\"Data augmentation disabled.\")\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3L9zSGtORKYI"
      },
      "source": [
        "\n",
        "## **3.3. Using weights from a pre-trained model as initial weights**\n",
        "---\n",
        "<font size = 4>  Here, you can set the the path to a pre-trained model from which the weights can be extracted and used as a starting point for this training session. **This pre-trained model needs to be a StarDist model**. \n",
        "\n",
        "<font size = 4> This option allows you to perform training over multiple Colab runtimes or to do transfer learning using models trained outside of ZeroCostDL4Mic. **You do not need to run this section if you want to train a network from scratch**.\n",
        "\n",
        "<font size = 4> In order to continue training from the point where the pre-trained model left off, it is adviseable to also **load the learning rate** that was used when the training ended. This is automatically saved for models trained with ZeroCostDL4Mic and will be loaded here. If no learning rate can be found in the model folder provided, the default learning rate will be used. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "9vC2n-HeLdiJ"
      },
      "outputs": [],
      "source": [
        "# @markdown ##Loading weights from a pre-trained network\n",
        "\n",
        "Use_pretrained_model = True #@param {type:\"boolean\"}\n",
        "\n",
        "pretrained_model_choice = \"Demo_3D_Model_from_Stardist_3D_paper\" #@param [\"Model_from_file\", \"Demo_3D_Model_from_Stardist_3D_paper\"]\n",
        "\n",
        "Weights_choice = \"best\" #@param [\"last\", \"best\"]\n",
        "\n",
        "\n",
        "#@markdown ###If you chose \"Model_from_file\", please provide the path to the model folder:\n",
        "pretrained_model_path = \"\" #@param {type:\"string\"}\n",
        "\n",
        "# --------------------- Check if we load a previously trained model ------------------------\n",
        "if Use_pretrained_model:\n",
        "\n",
        "# --------------------- Load the model from the choosen path ------------------------\n",
        "  if pretrained_model_choice == \"Model_from_file\":\n",
        "    h5_file_path = os.path.join(pretrained_model_path, \"weights_\"+Weights_choice+\".h5\")\n",
        "\n",
        "\n",
        "# --------------------- Download the Demo 3D model provided in the Stardist 3D github ------------------------\n",
        "\n",
        "  if pretrained_model_choice == \"Demo_3D_Model_from_Stardist_3D_paper\":\n",
        "    pretrained_model_name = \"Demo_3D\"\n",
        "    pretrained_model_path = \"/content/\"+pretrained_model_name\n",
        "    print(\"Downloading the Demo 3D model from the Stardist_3D paper\")\n",
        "    if os.path.exists(pretrained_model_path):\n",
        "      shutil.rmtree(pretrained_model_path)\n",
        "    os.makedirs(pretrained_model_path)\n",
        "    wget.download(\"https://raw.githubusercontent.com/mpicbg-csbd/stardist/master/models/examples/3D_demo/config.json\", pretrained_model_path)\n",
        "    wget.download(\"https://github.com/mpicbg-csbd/stardist/raw/master/models/examples/3D_demo/thresholds.json\", pretrained_model_path)\n",
        "    wget.download(\"https://github.com/mpicbg-csbd/stardist/blob/master/models/examples/3D_demo/weights_best.h5?raw=true\", pretrained_model_path)\n",
        "    wget.download(\"https://github.com/mpicbg-csbd/stardist/blob/master/models/examples/3D_demo/weights_last.h5?raw=true\", pretrained_model_path)\n",
        "    h5_file_path = os.path.join(pretrained_model_path, \"weights_\"+Weights_choice+\".h5\")\n",
        "\n",
        "# --------------------- Add additional pre-trained models here ------------------------\n",
        "\n",
        "\n",
        "\n",
        "# --------------------- Check the model exist ------------------------\n",
        "# If the model path chosen does not contain a pretrain model then use_pretrained_model is disabled, \n",
        "  if not os.path.exists(h5_file_path):\n",
        "    print(bcolors.WARNING+'WARNING: weights_last.h5 pretrained model does not exist'+W)\n",
        "    Use_pretrained_model = False\n",
        "\n",
        "  \n",
        "# If the model path contains a pretrain model, we load the training rate, \n",
        "  if os.path.exists(h5_file_path):\n",
        "#Here we check if the learning rate can be loaded from the quality control folder\n",
        "    if os.path.exists(os.path.join(pretrained_model_path, 'Quality Control', 'training_evaluation.csv')):\n",
        "\n",
        "      with open(os.path.join(pretrained_model_path, 'Quality Control', 'training_evaluation.csv'),'r') as csvfile:\n",
        "        csvRead = pd.read_csv(csvfile, sep=',')\n",
        "        #print(csvRead)\n",
        "    \n",
        "        if \"learning rate\" in csvRead.columns: #Here we check that the learning rate column exist (compatibility with model trained un ZeroCostDL4Mic bellow 1.4)\n",
        "          print(\"pretrained network learning rate found\")\n",
        "          #find the last learning rate\n",
        "          lastLearningRate = csvRead[\"learning rate\"].iloc[-1]\n",
        "          #Find the learning rate corresponding to the lowest validation loss\n",
        "          min_val_loss = csvRead[csvRead['val_loss'] == min(csvRead['val_loss'])]\n",
        "          #print(min_val_loss)\n",
        "          bestLearningRate = min_val_loss['learning rate'].iloc[-1]\n",
        "\n",
        "          if Weights_choice == \"last\":\n",
        "            print('Last learning rate: '+str(lastLearningRate))\n",
        "\n",
        "          if Weights_choice == \"best\":\n",
        "            print('Learning rate of best validation loss: '+str(bestLearningRate))\n",
        "\n",
        "        if not \"learning rate\" in csvRead.columns: #if the column does not exist, then initial learning rate is used instead\n",
        "          bestLearningRate = initial_learning_rate\n",
        "          lastLearningRate = initial_learning_rate\n",
        "          print(bcolors.WARNING+'WARNING: The learning rate cannot be identified from the pretrained network. Default learning rate of '+str(bestLearningRate)+' will be used instead' + W)\n",
        "\n",
        "#Compatibility with models trained outside ZeroCostDL4Mic but default learning rate will be used\n",
        "    if not os.path.exists(os.path.join(pretrained_model_path, 'Quality Control', 'training_evaluation.csv')):\n",
        "      print(bcolors.WARNING+'WARNING: The learning rate cannot be identified from the pretrained network. Default learning rate of '+str(initial_learning_rate)+' will be used instead'+ W)\n",
        "      bestLearningRate = initial_learning_rate\n",
        "      lastLearningRate = initial_learning_rate\n",
        "\n",
        "\n",
        "# Display info about the pretrained model to be loaded (or not)\n",
        "if Use_pretrained_model:\n",
        "  print(bcolors.WARNING+'Weights found in:')\n",
        "  print(h5_file_path)\n",
        "  print(bcolors.WARNING+'will be loaded prior to training.')\n",
        "\n",
        "else:\n",
        "  print(bcolors.WARNING+'No pretrained network will be used.')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MCGklf1vZf2M"
      },
      "source": [
        "#**4. Train the network**\n",
        "---"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1KYOuygETJkT"
      },
      "source": [
        "## **4.1. Prepare the training data and model for training**\n",
        "---\n",
        "<font size = 4>Here, we use the information from 3. to build the model and convert the training data into a suitable format for training."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "lIUAOJ_LMv5E"
      },
      "outputs": [],
      "source": [
        "#@markdown ##Create the model and dataset objects\n",
        "\n",
        "# --------------------- Here we delete the model folder if it already exist ------------------------\n",
        "\n",
        "if os.path.exists(model_path+'/'+model_name):\n",
        "  print(bcolors.WARNING +\"!! WARNING: Model folder already exists and has been removed !!\" + W)\n",
        "  shutil.rmtree(model_path+'/'+model_name)\n",
        "\n",
        "\n",
        "import warnings\n",
        "warnings.simplefilter(\"ignore\")\n",
        "\n",
        "# --------------------- Here we load the augmented data or the raw data ------------------------\n",
        "\n",
        "# if Use_Data_augmentation:\n",
        "#   Training_source_dir = Training_source_augmented\n",
        "#   Training_target_dir = Training_target_augmented\n",
        "\n",
        "# if not Use_Data_augmentation:\n",
        "#   Training_source_dir = Training_source\n",
        "#   Training_target_dir = Training_target\n",
        "# --------------------- ------------------------------------------------\n",
        "\n",
        "Training_source_dir = Training_source\n",
        "Training_target_dir = Training_target\n",
        "training_images_tiff=Training_source_dir+\"/*.tif\"\n",
        "mask_images_tiff=Training_target_dir+\"/*.tif\"\n",
        "\n",
        "\n",
        "# this funtion imports training images and masks and sorts them suitable for the network\n",
        "X = sorted(glob(training_images_tiff))   \n",
        "Y = sorted(glob(mask_images_tiff))   \n",
        "\n",
        "# assert -funtion check that X and Y really have images. If not this cell raises an error\n",
        "assert all(Path(x).name==Path(y).name for x,y in zip(X,Y))\n",
        "\n",
        "# Here we map the training dataset (images and masks).\n",
        "X = list(map(imread,X))\n",
        "Y = list(map(imread,Y))\n",
        "\n",
        "n_channel = 1 if X[0].ndim == 3 else X[0].shape[-1]\n",
        "\n",
        "\n",
        "\n",
        "#Normalize images and fill small label holes.\n",
        "axis_norm = (0,1,2)   # normalize channels independently\n",
        "# axis_norm = (0,1,2,3) # normalize channels jointly\n",
        "if n_channel > 1:\n",
        "    print(\"Normalizing image channels %s.\" % ('jointly' if axis_norm is None or 3 in axis_norm else 'independently'))\n",
        "    sys.stdout.flush()\n",
        "\n",
        "X = [normalize(x,1,99.8,axis=axis_norm) for x in tqdm(X)]\n",
        "Y = [fill_label_holes(y) for y in tqdm(Y)]\n",
        "\n",
        "#Here we split the your training dataset into training images (90 %) and validation images (10 %). \n",
        "\n",
        "assert len(X) > 1, \"not enough training data\"\n",
        "rng = np.random.RandomState(42)\n",
        "ind = rng.permutation(len(X))\n",
        "n_val = max(1, int(round(percentage * len(ind))))\n",
        "ind_train, ind_val = ind[:-n_val], ind[-n_val:]\n",
        "X_val, Y_val = [X[i] for i in ind_val]  , [Y[i] for i in ind_val]\n",
        "X_trn, Y_trn = [X[i] for i in ind_train], [Y[i] for i in ind_train] \n",
        "print('number of images: %3d' % len(X))\n",
        "print('- training:       %3d' % len(X_trn))\n",
        "print('- validation:     %3d' % len(X_val))\n",
        "\n",
        "\n",
        "extents = calculate_extents(Y)\n",
        "anisotropy = tuple(np.max(extents) / extents)\n",
        "print('empirical anisotropy of labeled objects = %s' % str(anisotropy))\n",
        "\n",
        "\n",
        "# Use OpenCL-based computations for data generator during training (requires 'gputools')\n",
        "use_gpu = False and gputools_available()\n",
        "\n",
        "\n",
        "#Here we ensure that our network has a minimal number of steps\n",
        "if (Use_Default_Advanced_Parameters) or (number_of_steps == 0):\n",
        "  number_of_steps = (Image_X//patch_size)*(Image_Y//patch_size)*(Image_Z//patch_height) * int(len(X)/batch_size)+1\n",
        "  if (Use_Data_augmentation):\n",
        "    number_of_steps = number_of_steps*6\n",
        "\n",
        "print(\"Number of steps: \"+str(number_of_steps))\n",
        "\n",
        "# --------------------- Using pretrained model ------------------------\n",
        "#Here we ensure that the learning rate set correctly when using pre-trained models\n",
        "if Use_pretrained_model:\n",
        "  if Weights_choice == \"last\":\n",
        "    initial_learning_rate = lastLearningRate\n",
        "\n",
        "  if Weights_choice == \"best\":            \n",
        "    initial_learning_rate = bestLearningRate\n",
        "# --------------------- ---------------------- ------------------------\n",
        "\n",
        "# Predict on subsampled grid for increased efficiency and larger field of view\n",
        "grid = tuple(1 if a > 1.5 else 2 for a in anisotropy)\n",
        "\n",
        "# Use rays on a Fibonacci lattice adjusted for measured anisotropy of the training data\n",
        "rays = Rays_GoldenSpiral(n_rays, anisotropy=anisotropy)\n",
        "\n",
        "conf = Config3D (\n",
        "    rays             = rays,\n",
        "    grid             = grid,\n",
        "    anisotropy       = anisotropy,\n",
        "    use_gpu          = use_gpu,\n",
        "    n_channel_in     = n_channel,\n",
        "    train_learning_rate = initial_learning_rate,\n",
        "    train_patch_size = (patch_height, patch_size, patch_size),\n",
        "    train_batch_size = batch_size,\n",
        ")\n",
        "print(conf)\n",
        "vars(conf)\n",
        "\n",
        "\n",
        "# --------------------- This is currently disabled as it give an error ------------------------\n",
        "#here we limit GPU to 80%\n",
        "if use_gpu:\n",
        "    from csbdeep.utils.tf import limit_gpu_memory\n",
        "    # adjust as necessary: limit GPU memory to be used by TensorFlow to leave some to OpenCL-based computations\n",
        "    limit_gpu_memory(0.8)\n",
        "# --------------------- ---------------------- ------------------------\n",
        "\n",
        "\n",
        "# Here we create a model according to section 5.3.\n",
        "model = StarDist3D(conf, name=model_name, basedir=trained_model)\n",
        "\n",
        "# --------------------- Using pretrained model ------------------------\n",
        "# Load the pretrained weights \n",
        "if Use_pretrained_model:\n",
        "  model.load_weights(h5_file_path)\n",
        "# --------------------- ---------------------- ------------------------\n",
        "\n",
        "\n",
        "#Here we check the FOV of the network.\n",
        "median_size = calculate_extents(Y, np.median)\n",
        "fov = np.array(model._axes_tile_overlap('ZYX'))\n",
        "if any(median_size > fov):\n",
        "    print(\"WARNING: median object size larger than field of view of the neural network.\")\n",
        "\n",
        "pdf_export(augmentation = Use_Data_augmentation, pretrained_model = Use_pretrained_model)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0Dfn8ZsEMv5d"
      },
      "source": [
        "## **4.2. Start Training**\n",
        "---\n",
        "\n",
        "<font size = 4>When playing the cell below you should see updates after each epoch (round). Network training can take some time.\n",
        "\n",
        "<font size = 4>* **CRITICAL NOTE:** Google Colab has a time limit for processing (to prevent using GPU power for datamining). Training time must be less than 12 hours! If training takes longer than 12 hours, please decrease the number of epochs or number of patches. Another way circumvent this is to save the parameters of the model after training and start training again from this point.\n",
        "\n",
        "<font size = 4>Once training is complete, the trained model is automatically saved on your Google Drive, in the **model_path** folder that was selected in Section 3. It is however wise to download the folder as all data can be erased at the next training if using the same folder."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "iwNmp1PUzRDQ",
        "scrolled": true
      },
      "outputs": [],
      "source": [
        "import time\n",
        "start = time.time()\n",
        "\n",
        "import warnings\n",
        "warnings.filterwarnings(\"ignore\")\n",
        "\n",
        "#@markdown ##Start training\n",
        "\n",
        "# augmenter = None\n",
        "\n",
        "# def augmenter(X_batch, Y_batch):\n",
        "#     \"\"\"Augmentation for data batch.\n",
        "#     X_batch is a list of input images (length at most batch_size)\n",
        "#     Y_batch is the corresponding list of ground-truth label images\n",
        "#     \"\"\"\n",
        "#     # ...\n",
        "#     return X_batch, Y_batch\n",
        "\n",
        "# Training the model. \n",
        "# 'input_epochs' and 'steps' refers to your input data in section 5.1 \n",
        "history = model.train(X_trn, Y_trn, validation_data=(X_val,Y_val), augmenter=augmenter,\n",
        "                      epochs=number_of_epochs, steps_per_epoch=number_of_steps)\n",
        "None;\n",
        "\n",
        "print(\"Training done\")\n",
        "\n",
        "\n",
        "# convert the history.history dict to a pandas DataFrame:     \n",
        "lossData = pd.DataFrame(history.history) \n",
        "\n",
        "if os.path.exists(model_path+\"/\"+model_name+\"/Quality Control\"):\n",
        "  shutil.rmtree(model_path+\"/\"+model_name+\"/Quality Control\")\n",
        "\n",
        "os.makedirs(model_path+\"/\"+model_name+\"/Quality Control\")\n",
        "\n",
        "# The training evaluation.csv is saved (overwrites the Files if needed). \n",
        "lossDataCSVpath = model_path+'/'+model_name+'/Quality Control/training_evaluation.csv'\n",
        "with open(lossDataCSVpath, 'w') as f:\n",
        "  writer = csv.writer(f)\n",
        "  writer.writerow(['loss','val_loss', 'learning rate'])\n",
        "  for i in range(len(history.history['loss'])):\n",
        "    writer.writerow([history.history['loss'][i], history.history['val_loss'][i], history.history['lr'][i]])\n",
        "\n",
        "\n",
        "print(\"Network optimization in progress\")\n",
        "\n",
        "#Here we optimize the network.\n",
        "model.optimize_thresholds(X_val, Y_val)\n",
        "print(\"Done\")\n",
        "\n",
        "\n",
        "# Displaying the time elapsed for training\n",
        "dt = time.time() - start\n",
        "mins, sec = divmod(dt, 60) \n",
        "hour, mins = divmod(mins, 60) \n",
        "print(\"Time elapsed:\",hour, \"hour(s)\",mins,\"min(s)\",round(sec),\"sec(s)\")\n",
        "\n",
        "#Create a pdf document with training summary\n",
        "pdf_export(trained = True, augmentation = Use_Data_augmentation, pretrained_model = Use_pretrained_model)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_0Hynw3-xHp1"
      },
      "source": [
        "# **5. Evaluate your model**\n",
        "---\n",
        "\n",
        "<font size = 4>This section allows the user to perform important quality checks on the validity and generalisability of the trained model. \n",
        "\n",
        "<font size = 4>**We highly recommend to perform quality control on all newly trained models.**\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "eAJzMwPA6tlH"
      },
      "outputs": [],
      "source": [
        "# model name and path\n",
        "#@markdown ###Do you want to assess the model you just trained ?\n",
        "Use_the_current_trained_model = True #@param {type:\"boolean\"}\n",
        "\n",
        "#@markdown ###If not, please provide the path to the model folder:\n",
        "\n",
        "QC_model_folder = \"\" #@param {type:\"string\"}\n",
        "\n",
        "#Here we define the loaded model name and path\n",
        "QC_model_name = os.path.basename(QC_model_folder)\n",
        "QC_model_path = os.path.dirname(QC_model_folder)\n",
        "\n",
        "\n",
        "\n",
        "if (Use_the_current_trained_model): \n",
        "  QC_model_name = model_name\n",
        "  QC_model_path = model_path\n",
        "\n",
        "full_QC_model_path = QC_model_path+'/'+QC_model_name+'/'\n",
        "if os.path.exists(full_QC_model_path):\n",
        "  print(\"The \"+QC_model_name+\" network will be evaluated\")\n",
        "else:\n",
        "  W  = '\\033[0m'  # white (normal)\n",
        "  R  = '\\033[31m' # red\n",
        "  print(R+'!! WARNING: The chosen model does not exist !!'+W)\n",
        "  print('Please make sure you provide a valid model path and model name before proceeding further.')\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dhJROwlAMv5o"
      },
      "source": [
        "## **5.1. Inspection of the loss function**\n",
        "---\n",
        "\n",
        "<font size = 4>First, it is good practice to evaluate the training progress by comparing the training loss with the validation loss. The latter is a metric which shows how well the network performs on a subset of unseen data which is set aside from the training dataset. For more information on this, see for example [this review](https://www.ncbi.nlm.nih.gov/pmc/articles/PMC6381354/) by Nichols *et al.*\n",
        "\n",
        "<font size = 4>**Training loss** describes an error value after each epoch for the difference between the model's prediction and its ground-truth target.\n",
        "\n",
        "<font size = 4>**Validation loss** describes the same error value between the model's prediction on a validation image and compared to it's target.\n",
        "\n",
        "<font size = 4>During training both values should decrease before reaching a minimal value which does not decrease further even after more training. Comparing the development of the validation loss with the training loss can give insights into the model's performance.\n",
        "\n",
        "<font size = 4>Decreasing **Training loss** and **Validation loss** indicates that training is still necessary and increasing the `number_of_epochs` is recommended. Note that the curves can look flat towards the right side, just because of the y-axis scaling. The network has reached convergence once the curves flatten out. After this point no further training is required. If the **Validation loss** suddenly increases again an the **Training loss** simultaneously goes towards zero, it means that the network is overfitting to the training data. In other words the network is remembering the exact patterns from the training data and no longer generalizes well to unseen data. In this case the training dataset has to be increased."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "vMzSP50kMv5p"
      },
      "outputs": [],
      "source": [
        "#@markdown ##Play the cell to show a plot of training errors vs. epoch number\n",
        "import csv\n",
        "from matplotlib import pyplot as plt\n",
        "\n",
        "lossDataFromCSV = []\n",
        "vallossDataFromCSV = []\n",
        "\n",
        "with open(QC_model_path+'/'+QC_model_name+'/Quality Control/training_evaluation.csv','r') as csvfile:\n",
        "    csvRead = csv.reader(csvfile, delimiter=',')\n",
        "    next(csvRead)\n",
        "    for row in csvRead:\n",
        "        lossDataFromCSV.append(float(row[0]))\n",
        "        vallossDataFromCSV.append(float(row[1]))\n",
        "\n",
        "epochNumber = range(len(lossDataFromCSV))\n",
        "plt.figure(figsize=(15,10))\n",
        "\n",
        "plt.subplot(2,1,1)\n",
        "plt.plot(epochNumber,lossDataFromCSV, label='Training loss')\n",
        "plt.plot(epochNumber,vallossDataFromCSV, label='Validation loss')\n",
        "plt.title('Training loss and validation loss vs. epoch number (linear scale)')\n",
        "plt.ylabel('Loss')\n",
        "plt.xlabel('Epoch number')\n",
        "plt.legend()\n",
        "\n",
        "\n",
        "plt.subplot(2,1,2)\n",
        "plt.semilogy(epochNumber,lossDataFromCSV, label='Training loss')\n",
        "plt.semilogy(epochNumber,vallossDataFromCSV, label='Validation loss')\n",
        "plt.title('Training loss and validation loss vs. epoch number (log scale)')\n",
        "plt.ylabel('Loss')\n",
        "plt.xlabel('Epoch number')\n",
        "plt.legend()\n",
        "plt.savefig(QC_model_path+'/'+QC_model_name+'/Quality Control/lossCurvePlots.png',bbox_inches='tight',pad_inches=0)\n",
        "plt.show()\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "X5_92nL2xdP6"
      },
      "source": [
        "## **5.2. Error mapping and quality metrics estimation**\n",
        "---\n",
        "<font size = 4>This section will calculate the Intersection over Union score for all the images provided in the Source_QC_folder and Target_QC_folder ! The result for one of the image will also be displayed.\n",
        "\n",
        "<font size = 4>The **Intersection over Union** metric is a method that can be used to quantify the percent overlap between the target mask and your prediction output. **Therefore, the closer to 1, the better the performance.** This metric can be used to assess the quality of your model to accurately predict nuclei. \n",
        "\n",
        "<font size = 4> The results can be found in the \"*Quality Control*\" folder which is located inside your \"model_folder\"."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "w90MdriMxhjD"
      },
      "outputs": [],
      "source": [
        "#@markdown ##Give the paths to an image to test the performance of the model with.\n",
        "\n",
        "import warnings\n",
        "warnings.filterwarnings(\"ignore\")\n",
        "\n",
        "Source_QC_folder = \"\" #@param{type:\"string\"}\n",
        "Target_QC_folder = \"\" #@param{type:\"string\"}\n",
        "\n",
        "#Here we allow the user to choose the number of tile to be used when predicting the images\n",
        "#@markdown #####To analyse large image, your images need to be divided into tiles.  Each tile will then be processed independently and re-assembled to generate the final image. \"Automatic_number_of_tiles\" will search for and use the smallest number of tiles that can be used, at the expanse of your runtime. Alternatively, manually input the number of tiles in each dimension to be used to process your images. \n",
        "\n",
        "Automatic_number_of_tiles = False #@param {type:\"boolean\"}\n",
        "#@markdown #####If you get an Out of memory (OOM) error when using the \"Automatic_number_of_tiles\" option, disable it and manually input the values to be used to process your images.  Progressively increases these numbers until the OOM error disappear.\n",
        "n_tiles_Z =  1#@param {type:\"number\"}\n",
        "n_tiles_Y =  1#@param {type:\"number\"}\n",
        "n_tiles_X =  1#@param {type:\"number\"}\n",
        "\n",
        "if (Automatic_number_of_tiles): \n",
        "  n_tilesZYX = None\n",
        "\n",
        "if not (Automatic_number_of_tiles):\n",
        "  n_tilesZYX = (n_tiles_Z, n_tiles_Y, n_tiles_X)\n",
        "\n",
        "\n",
        "#Create a quality control Folder and check if the folder already exist\n",
        "if os.path.exists(QC_model_path+\"/\"+QC_model_name+\"/Quality Control\") == False:\n",
        "  os.makedirs(QC_model_path+\"/\"+QC_model_name+\"/Quality Control\")\n",
        "\n",
        "if os.path.exists(QC_model_path+\"/\"+QC_model_name+\"/Quality Control/Prediction\"):\n",
        "  shutil.rmtree(QC_model_path+\"/\"+QC_model_name+\"/Quality Control/Prediction\")\n",
        "\n",
        "os.makedirs(QC_model_path+\"/\"+QC_model_name+\"/Quality Control/Prediction\")\n",
        "\n",
        "\n",
        "# Generate predictions from the Source_QC_folder and save them in the QC folder\n",
        "\n",
        "Source_QC_folder_tif = Source_QC_folder+\"/*.tif\"\n",
        "\n",
        "\n",
        "np.random.seed(16)\n",
        "lbl_cmap = random_label_cmap()\n",
        "Z = sorted(glob(Source_QC_folder_tif))\n",
        "Z = list(map(imread,Z))\n",
        "n_channel = 1 if Z[0].ndim == 2 else Z[0].shape[-1]\n",
        "axis_norm = (0,1)   # normalize channels independently\n",
        "\n",
        "print('Number of test dataset found in the folder: '+str(len(Z)))\n",
        "\n",
        "  \n",
        "  # axis_norm = (0,1,2) # normalize channels jointly\n",
        "if n_channel > 1:\n",
        "  print(\"Normalizing image channels %s.\" % ('jointly' if axis_norm is None or 2 in axis_norm else 'independently'))\n",
        "\n",
        "model = StarDist3D(None, name=QC_model_name, basedir=QC_model_path)\n",
        "\n",
        "names = [os.path.basename(f) for f in sorted(glob(Source_QC_folder_tif))]\n",
        "\n",
        " \n",
        "# modify the names to suitable form: path_images/image_numberX.tif\n",
        "  \n",
        "lenght_of_Z = len(Z)\n",
        "  \n",
        "for i in range(lenght_of_Z):\n",
        "  img = normalize(Z[i], 1,99.8, axis=axis_norm)\n",
        "  labels, polygons = model.predict_instances(img, n_tiles=n_tilesZYX)\n",
        "  os.chdir(QC_model_path+\"/\"+QC_model_name+\"/Quality Control/Prediction\")\n",
        "  #imsave(names[i], labels, polygons)\n",
        "  save_tiff_imagej_compatible(names[i], labels, axes='ZYX')\n",
        "\n",
        "\n",
        "# Here we start testing the differences between GT and predicted masks\n",
        "\n",
        "\n",
        "with open(QC_model_path+\"/\"+QC_model_name+\"/Quality Control/Quality_Control for \"+QC_model_name+\".csv\", \"w\", newline='') as file:\n",
        "  writer = csv.writer(file)\n",
        "  writer.writerow([\"image\",\"Prediction v. GT Intersection over Union\"])  \n",
        "\n",
        "  # Initialise the lists \n",
        "  filename_list = []\n",
        "  IoU_score_list = []\n",
        "\n",
        "  for n in os.listdir(Source_QC_folder):\n",
        "    if not os.path.isdir(os.path.join(Source_QC_folder,n)):\n",
        "      print('Running QC on: '+n)\n",
        "    \n",
        "      test_input = io.imread(os.path.join(Source_QC_folder,n))\n",
        "      test_prediction = io.imread(os.path.join(QC_model_path+\"/\"+QC_model_name+\"/Quality Control/Prediction\",n))\n",
        "      test_ground_truth_image = io.imread(os.path.join(Target_QC_folder, n))\n",
        "\n",
        "#Convert pixel values to 0 or 255\n",
        "      test_prediction_0_to_255 = test_prediction\n",
        "      test_prediction_0_to_255[test_prediction_0_to_255>0] = 255\n",
        "\n",
        "#Convert pixel values to 0 or 255\n",
        "      test_ground_truth_0_to_255 = test_ground_truth_image\n",
        "      test_ground_truth_0_to_255[test_ground_truth_0_to_255>0] = 255\n",
        "\n",
        "# Intersection over Union metric\n",
        "\n",
        "      intersection = np.logical_and(test_ground_truth_0_to_255, test_prediction_0_to_255)\n",
        "      union = np.logical_or(test_ground_truth_0_to_255, test_prediction_0_to_255)\n",
        "      iou_score =  np.sum(intersection) / np.sum(union)\n",
        "      writer.writerow([n, str(iou_score)])\n",
        "\n",
        "      print(\"IoU: \"+str(round(iou_score,3)))\n",
        "\n",
        "      filename_list.append(n)\n",
        "      IoU_score_list.append(iou_score)\n",
        "\n",
        "\n",
        "\n",
        "# Table with metrics as dataframe output\n",
        "pdResults = pd.DataFrame(index = filename_list)\n",
        "pdResults[\"IoU\"] = IoU_score_list\n",
        "\n",
        "# Display results\n",
        "pdResults.head()\n",
        "\n",
        "\n",
        "# ------------- For display ------------\n",
        "print('--------------------------------------------------------------')\n",
        "@interact\n",
        "def show_QC_results(file=os.listdir(Source_QC_folder)):\n",
        "\n",
        "  f=plt.figure(figsize=(32,8))\n",
        "\n",
        "  test_input = io.imread(os.path.join(Source_QC_folder, file))\n",
        "  test_prediction = io.imread(os.path.join(QC_model_path+\"/\"+QC_model_name+\"/Quality Control/Prediction\", file))\n",
        "  test_ground_truth_image = io.imread(os.path.join(Target_QC_folder, file))\n",
        "\n",
        "  norm = simple_norm(test_input, percent = 99)\n",
        "  Image_Z = test_input.shape[0]\n",
        "  mid_plane = int(Image_Z / 2)+1\n",
        "\n",
        "  #Convert pixel values to 0 or 255\n",
        "  test_prediction_0_to_255 = test_prediction\n",
        "  test_prediction_0_to_255[test_prediction_0_to_255>0] = 255\n",
        "\n",
        "  #Convert pixel values to 0 or 255\n",
        "  test_ground_truth_0_to_255 = test_ground_truth_image\n",
        "  test_ground_truth_0_to_255[test_ground_truth_0_to_255>0] = 255\n",
        "\n",
        "  #Input\n",
        "  plt.subplot(1,4,1)\n",
        "  plt.axis('off')\n",
        "  plt.imshow(test_input[mid_plane], aspect='equal', norm=norm, cmap='magma',   interpolation='nearest')\n",
        "  plt.title('Input')\n",
        "\n",
        "  #Ground-truth\n",
        "  plt.subplot(1,4,2)\n",
        "  plt.axis('off')\n",
        "  plt.imshow(test_ground_truth_0_to_255[mid_plane], aspect='equal', cmap='Greens')\n",
        "  plt.title('Ground Truth')\n",
        "\n",
        "  #Prediction\n",
        "  plt.subplot(1,4,3)\n",
        "  plt.axis('off')\n",
        "  plt.imshow(test_prediction_0_to_255[mid_plane], aspect='equal', cmap='Purples')\n",
        "  plt.title('Prediction')\n",
        "\n",
        "  #Overlay\n",
        "  plt.subplot(1,4,4)\n",
        "  plt.axis('off')\n",
        "  plt.imshow(test_ground_truth_0_to_255[mid_plane], cmap='Greens')\n",
        "  plt.imshow(test_prediction_0_to_255[mid_plane], alpha=0.5, cmap='Purples')\n",
        "  plt.title('Ground Truth and Prediction, Intersection over Union:'+str(round(pdResults.loc[file][\"IoU\"],3)))\n",
        "  plt.savefig(full_QC_model_path+'/Quality Control/QC_example_data.png',bbox_inches='tight',pad_inches=0)\n",
        "\n",
        " \n",
        "\n",
        "# Make a pdf summary of the QC results\n",
        "qc_pdf_export()\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-tJeeJjLnRkP"
      },
      "source": [
        "# **6. Using the trained model**\n",
        "\n",
        "---\n",
        "\n",
        "<font size = 4>In this section the unseen data is processed using the trained model (in section 4). First, your unseen images are uploaded and prepared for prediction. After that your trained model from section 4 is activated and finally saved into your Google Drive."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d8wuQGjoq6eN"
      },
      "source": [
        "## **6.1. Generate prediction(s) from unseen dataset**\n",
        "---\n",
        "\n",
        "<font size = 4>The current trained model (from section 4.3) can now be used to process images. If an older model needs to be used, please untick the **Use_the_current_trained_model** box and enter the name and path of the model to use. Predicted output images are saved in your **Prediction_folder** folder as restored image stacks (ImageJ-compatible TIFF images).\n",
        "\n",
        "<font size = 4>**`Data_folder`:** This folder should contains the images that you want to predict using the network that you trained.\n",
        "\n",
        "<font size = 4>**`Result_folder`:** This folder will contain the predicted output ROI.\n",
        "\n",
        "<font size = 4>**`Data_type`:** Please indicate if the images you want to predict are single images or stacks\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "y2TD5p7MZrEb"
      },
      "outputs": [],
      "source": [
        "from PIL import Image\n",
        "\n",
        "#@markdown ### Provide the path to your dataset and to the folder where the prediction will be saved (Result folder), then play the cell to predict output on your unseen images.\n",
        "\n",
        "Data_folder = \"\" #@param {type:\"string\"}\n",
        "#test_dataset = Data_folder\n",
        "\n",
        "Results_folder = \"\" #@param {type:\"string\"}\n",
        "#results = results_folder\n",
        "\n",
        "\n",
        "# model name and path\n",
        "#@markdown ###Do you want to use the current trained model?\n",
        "Use_the_current_trained_model = True #@param {type:\"boolean\"}\n",
        "\n",
        "\n",
        "#@markdown ###If not, please provide the path to the model folder:\n",
        "\n",
        "Prediction_model_folder = \"\" #@param {type:\"string\"}\n",
        "\n",
        "#Here we find the loaded model name and parent path\n",
        "Prediction_model_name = os.path.basename(Prediction_model_folder)\n",
        "Prediction_model_path = os.path.dirname(Prediction_model_folder)\n",
        "\n",
        "#Here we allow the user to choose the number of tile to be used when predicting the images\n",
        "#@markdown #####To analyse large image, your images need to be divided into tiles.  Each tile will then be processed independently and re-assembled to generate the final image. \"Automatic_number_of_tiles\" will search for and use the smallest number of tiles that can be used, at the expanse of your runtime. Alternatively, manually input the number of tiles in each dimension to be used to process your images. \n",
        "\n",
        "Automatic_number_of_tiles = False #@param {type:\"boolean\"}\n",
        "#@markdown #####If you get an Out of memory (OOM) error when using the \"Automatic_number_of_tiles\" option, disable it and manually input the values to be used to process your images.  Progressively increases these numbers until the OOM error disappear.\n",
        "n_tiles_Z =  1#@param {type:\"number\"}\n",
        "n_tiles_Y =  1#@param {type:\"number\"}\n",
        "n_tiles_X =  1#@param {type:\"number\"}\n",
        "\n",
        "if (Automatic_number_of_tiles): \n",
        "  n_tilesZYX = None\n",
        "\n",
        "if not (Automatic_number_of_tiles):\n",
        "  n_tilesZYX = (n_tiles_Z, n_tiles_Y, n_tiles_X)\n",
        "\n",
        "\n",
        "if (Use_the_current_trained_model): \n",
        "  print(\"Using current trained network\")\n",
        "  Prediction_model_name = model_name\n",
        "  Prediction_model_path = model_path\n",
        "\n",
        "full_Prediction_model_path = Prediction_model_path+'/'+Prediction_model_name+'/'\n",
        "if os.path.exists(full_Prediction_model_path):\n",
        "  print(\"The \"+Prediction_model_name+\" network will be used.\")\n",
        "else:\n",
        "  W  = '\\033[0m'  # white (normal)\n",
        "  R  = '\\033[31m' # red\n",
        "  print(R+'!! WARNING: The chosen model does not exist !!'+W)\n",
        "  print('Please make sure you provide a valid model path and model name before proceeding further.')\n",
        "\n",
        "\n",
        "#single images\n",
        "#testDATA = test_dataset\n",
        "Dataset = Data_folder+\"/*.tif\"\n",
        "\n",
        "\n",
        "np.random.seed(16)\n",
        "lbl_cmap = random_label_cmap()\n",
        "X = sorted(glob(Dataset))\n",
        "X = list(map(imread,X))\n",
        "n_channel = 1 if X[0].ndim == 2 else X[0].shape[-1]\n",
        "axis_norm = (0,1)   # normalize channels independently\n",
        "  \n",
        "# axis_norm = (0,1,2) # normalize channels jointly\n",
        "if n_channel > 1:\n",
        "  print(\"Normalizing image channels %s.\" % ('jointly' if axis_norm is None or 2 in axis_norm else 'independently'))\n",
        "model = StarDist3D(None, name=Prediction_model_name, basedir=Prediction_model_path)\n",
        " \n",
        "#Sorting and mapping original test dataset\n",
        "#X = sorted(glob(Dataset))\n",
        "#X = list(map(imread,X))\n",
        "names = [os.path.basename(f) for f in sorted(glob(Dataset))]\n",
        "\n",
        "# modify the names to suitable form: path_images/image_numberX.tif\n",
        "FILEnames=[]\n",
        "for m in names:\n",
        "  m=Results_folder+'/'+m\n",
        "  FILEnames.append(m)\n",
        "\n",
        "  # Predictions folder\n",
        "lenght_of_X = len(X)\n",
        "for i in range(lenght_of_X):\n",
        "  img = normalize(X[i], 1,99.8, axis=axis_norm)\n",
        "  labels, polygons = model.predict_instances(img, n_tiles=n_tilesZYX)\n",
        "  \n",
        "# Save the predicted mask in the result folder\n",
        "  os.chdir(Results_folder)\n",
        "  #imsave(FILEnames[i], labels, polygons)\n",
        "  save_tiff_imagej_compatible(FILEnames[i], labels, axes='ZYX')\n",
        "\n",
        "\n",
        "\n",
        "print(\"The mid-plane image is displayed below.\")\n",
        "# ------------- For display ------------\n",
        "print('--------------------------------------------------------------')\n",
        "@interact\n",
        "def show_QC_results(file=os.listdir(Data_folder)):\n",
        "  plt.figure(figsize=(13,10))\n",
        "\n",
        "  img = imread(os.path.join(Data_folder, file))\n",
        "  img = normalize(img, 1,99.8, axis=axis_norm)\n",
        "  labels = imread(os.path.join(Results_folder, file))\n",
        "  z = max(0, img.shape[0] // 2 - 5)\n",
        "\n",
        "  plt.subplot(121)\n",
        "  plt.imshow((img if img.ndim==3 else img[...,:3])[z], clim=(0,1), cmap='gray')\n",
        "  plt.title('Raw image (XY slice)')\n",
        "  plt.axis('off')\n",
        "  plt.subplot(122)\n",
        "  plt.imshow((img if img.ndim==3 else img[...,:3])[z], clim=(0,1), cmap='gray')\n",
        "  plt.imshow(labels[z], cmap=lbl_cmap, alpha=0.5)\n",
        "  plt.title('Image and predicted labels (XY slice)')\n",
        "  plt.axis('off');\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hvkd66PldsXB"
      },
      "source": [
        "## **6.2. Download your predictions**\n",
        "---\n",
        "\n",
        "<font size = 4>**Store your data** and ALL its results elsewhere by downloading it from Google Drive and after that clean the original folder tree (datasets, results, trained model etc.) if you plan to train or use new networks. Please note that the notebook will otherwise **OVERWRITE** all files which have the same name."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4DSziojsE6PC"
      },
      "source": [
        "# **7. Version log**\n",
        "---\n",
        "\n",
        "<font size = 4>**v1.15.2**:  \n",
        "\n",
        "*  The way that the label are saved has been updated\n",
        "\n",
        "<font size = 4>**v1.15.1**:  \n",
        "\n",
        "*  Tensorflow now v2.8\n",
        "*  `__future__` dependency is imported at the beginning\n",
        "\n",
        "<font size = 4>**v1.13**:  \n",
        "\n",
        "*  StarDist is now downgraded to v 0.6.2 to ensure compatibility with previously trained models.\n",
        "\n",
        "*   This version now includes an automatic restart allowing to set the h5py library to v2.10.\n",
        "*  The section 1 and 2 are now swapped for better export of *requirements.txt*.\n",
        "This version also now includes built-in version check and the version log that \n",
        "\n",
        "*   This version also now includes built-in version check and the version log that you're reading now.\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UvSlTaH14s3t"
      },
      "source": [
        "#**Thank you for using StarDist 3D!**"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "machine_shape": "hm",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "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.6.7"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}