Colab_notebooks/StarDist_2D_ZeroCostDL4Mic.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IkSguVy8Xv83"
      },
      "source": [
        "# **StarDist (2D)**\n",
        "---\n",
        "\n",
        "<font size = 4>**StarDist 2D** is a deep-learning method that can be used to segment cell nuclei from bioimages and was first published by [Schmidt *et al.* in 2018, on arXiv](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 2D network is based on an adapted U-Net network architecture.\n",
        "\n",
        "<font size = 4> **This particular notebook enables nuclei segmentation of 2D dataset. If you are interested in 3D dataset, you should use the StarDist 3D 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. The BioImage Model Zoo export was jointly developed by [Estibaliz Gómez de Mariscal](https://github.com/esgomezm) (deepImageJ team).\n",
        "\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>The guidelines to use the trained network in ImageJ with deepImageJ are given in the following paper:\n",
        "\n",
        "<font size = 4>**DeepImageJ: a user-friendly environment to run deep learning models in ImageJ**, bioRxiv (2019) by *Estibaliz Gómez-de-Mariscal, Carlos García-López-de-Haro, Wei Ouyang, Laurène Donati, Emma Lundberg, Michael Unser, Arrate Muñoz-Barrutia and Daniel Sage* (https://doi.org/10.1101/799270)\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**. 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. This can include Test dataset for which you have the equivalent output and can compare to what the network provides.\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": "XuwTHSva_Y5K"
      },
      "source": [
        "## **1.1. Install key dependencies**\n",
        "---\n",
        "<font size = 4> "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "fq21zJVFNASx"
      },
      "outputs": [],
      "source": [
        "#@markdown ##Install StarDist and dependencies\n",
        "\n",
        "\n",
        "# Install packages required for image processing and analysis\n",
        "\n",
        "# Package for handling tiff-files\n",
        "!pip install -q tifffile\n",
        "\n",
        "# Package for restoration of fluorescence microscopy images using deep learning (CARE)\n",
        "# It relies on Keras and TensorFlow\n",
        "!pip install -q csbdeep\n",
        "\n",
        "# Package for working with STARDIST, a method for object detection and segmentation in images\n",
        "!pip install -q stardist\n",
        "\n",
        "# The [bioimageio] extras are installed for bioimageio model support\n",
        "!pip install 'bioimageio.core>=0.5,<0.6'\n",
        "\n",
        "# Package for improving STARDIST performance using GPU acceleration\n",
        "!pip install -q gputools\n",
        "\n",
        "# Package for downloading files from the web\n",
        "!pip install -q wget\n",
        "\n",
        "# Package for generating PDF documents\n",
        "!pip install -q fpdf2\n",
        "\n",
        "# Package for creating nice-looking tables\n",
        "!pip install -q PTable\n",
        "\n",
        "# Package for working with zarr, a format for efficient storage of large arrays\n",
        "!pip install -q zarr\n",
        "\n",
        "# Package for image compression and decompression using various codecs\n",
        "!pip install -q imagecodecs\n",
        "\n",
        "# Overall functionality: This code installs various packages necessary for image processing,\n",
        "# including handling image files, restoration of fluorescence microscopy images,\n",
        "# object detection and segmentation, GPU acceleration, file downloading, PDF generation,\n",
        "# table creation, storage of large arrays, and image compression.\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "#Force session restart\n",
        "#exit(0)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e_oQT-9180CX"
      },
      "source": [
        "\n",
        "## **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",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2uHigbVJ9CUh"
      },
      "source": [
        "## **1.3. Load key dependencies**\n",
        "---\n",
        "<font size = 4> "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "HMAdm-Mc9HFz"
      },
      "outputs": [],
      "source": [
        "from __future__ import print_function, unicode_literals, absolute_import, division\n",
        "#@markdown ##Load key dependencies\n",
        "Notebook_version = '1.19.3'\n",
        "Network = 'StarDist 2D'\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",
        "    print(filtered_list)\n",
        "\n",
        "    file=open(path,'w')\n",
        "    for item in filtered_list:\n",
        "      print(item)\n",
        "      file.writelines(item)\n",
        "\n",
        "    file.close()\n",
        "\n",
        "import sys\n",
        "before = [str(m) for m in sys.modules]\n",
        "\n",
        "\n",
        "#%load_ext memory_profiler\n",
        "import tensorflow\n",
        "print(tensorflow.__version__)\n",
        "print(\"Tensorflow enabled.\")\n",
        "import imagecodecs\n",
        "\n",
        "# ------- Variable specific to Stardist -------\n",
        "from stardist import fill_label_holes, random_label_cmap, calculate_extents, gputools_available, relabel_image_stardist, random_label_cmap,  relabel_image_stardist, _draw_polygons, export_imagej_rois\n",
        "from stardist.models import Config2D, StarDist2D, StarDistData2D # import objects\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",
        "%matplotlib inline\n",
        "%config InlineBackend.figure_format = 'retina'\n",
        "\n",
        "from PIL import Image\n",
        "import zarr\n",
        "from zipfile import ZIP_DEFLATED\n",
        "from csbdeep.data import Normalizer, normalize_mi_ma\n",
        "import imagecodecs\n",
        "\n",
        "\n",
        "class MyNormalizer(Normalizer):\n",
        "    def __init__(self, mi, ma):\n",
        "            self.mi, self.ma = mi, ma\n",
        "    def before(self, x, axes):\n",
        "        return normalize_mi_ma(x, self.mi, self.ma, dtype=np.float32)\n",
        "    def after(*args, **kwargs):\n",
        "        assert False\n",
        "    @property\n",
        "    def do_after(self):\n",
        "        return False\n",
        "\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, imwrite\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, img_as_ubyte, img_as_float\n",
        "from skimage.util import img_as_ubyte\n",
        "from tqdm import tqdm \n",
        "import cv2\n",
        "from fpdf import FPDF, HTMLMixin\n",
        "from datetime import datetime\n",
        "from pip._internal.operations.freeze import freeze\n",
        "import subprocess\n",
        "\n",
        "# For sliders and dropdown menu and progress bar\n",
        "from ipywidgets import interact\n",
        "import ipywidgets as widgets\n",
        "\n",
        "import urllib\n",
        "from tensorflow.keras.layers import Concatenate\n",
        "from tensorflow.keras.models import Model\n",
        "\n",
        "# BioImage Model Zoo export\n",
        "from stardist import export_bioimageio, import_bioimageio\n",
        "from pathlib import Path\n",
        "import bioimageio.core\n",
        "from bioimageio.core import load_resource_description\n",
        "\n",
        "#Create a variable to get and store relative base path\n",
        "base_path = os.getcwd()\n",
        "\n",
        "# Colors for the warning messages\n",
        "class bcolors:\n",
        "  WARNING = '\\033[31m'\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('------------------------------------------')\n",
        "print(\"Libraries installed\")\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",
        "# PDF export\n",
        "\n",
        "def pdf_export(trained=False, augmentation = False, pretrained_model = False):\n",
        "  class MyFPDF(FPDF, HTMLMixin):\n",
        "    pass\n",
        "  from skimage import io\n",
        "  pdf = MyFPDF()\n",
        "  pdf.add_page()\n",
        "  pdf.set_right_margin(-1)\n",
        "  pdf.set_font(\"Arial\", size = 11, style='B') \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 method:'\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",
        "\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",
        "  try:\n",
        "    cuda_version = subprocess.run([\"nvcc\",\"--version\"],stdout=subprocess.PIPE)\n",
        "    cuda_version = cuda_version.stdout.decode('utf-8')\n",
        "    cuda_version = cuda_version[cuda_version.find(', V')+3:-1]\n",
        "  except:\n",
        "    cuda_version = ' - No cuda found - '\n",
        "  try:\n",
        "    gpu_name = subprocess.run([\"nvidia-smi\"],stdout=subprocess.PIPE)\n",
        "    gpu_name = gpu_name.stdout.decode('utf-8')\n",
        "    gpu_name = gpu_name[gpu_name.find('Tesla'):gpu_name.find('Tesla')+10]\n",
        "  except:\n",
        "    gpu_name = ' - No GPU found - '\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], as_gray = True).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(dataset_size)+' paired image patches (image dimensions: '+str(shape)+', patch size: ('+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+') (von Chamier & Laine et al., 2021). Key python packages used include tensorflow (v '+version_numbers[0]+'), csbdeep (v '+version_numbers[3]+'), 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(dataset_size)+' paired image patches (image dimensions: '+str(shape)+', patch size: ('+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+') (von Chamier & Laine et al., 2021). The model was retrained from a pretrained model. Key python packages used include tensorflow (v '+version_numbers[0]+'), csbdeep (v '+version_numbers[3]+'), 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 by a factor of '+str(Multiply_dataset_by)\n",
        "    \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>number_of_epochs</td>\n",
        "      <td>{0}</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <td>patch_size</td>\n",
        "      <td>{1}</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <td>batch_size</td>\n",
        "      <td>{2}</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <td>number_of_steps</td>\n",
        "      <td>{3}</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <td>percentage_validation</td>\n",
        "      <td>{4}</td>\n",
        "    </tr>\n",
        "      <tr>\n",
        "      <td>n_rays</td>\n",
        "      <td>{5}</td>\n",
        "    </tr>\n",
        "      <tr>\n",
        "      <td>grid_parameter</td>\n",
        "      <td>{6}</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <td>initial_learning_rate</td>\n",
        "      <td>{7}</td>\n",
        "    </tr>\n",
        "  </table>\n",
        "  \"\"\".format(number_of_epochs,str(patch_size)+'x'+str(patch_size),batch_size,number_of_steps,percentage_validation,n_rays,grid_parameter,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(base_path + '/TrainingDataExample_StarDist2D.png').shape\n",
        "  pdf.image(base_path + '/TrainingDataExample_StarDist2D.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 2D: 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",
        "  if 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",
        "  from skimage import io\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 2D'\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",
        "\n",
        "  <table width=100% 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",
        "    fp = header[2]\n",
        "    tp = header[3]\n",
        "    fn = header[4]\n",
        "    precision = header[5]\n",
        "    recall = header[6]\n",
        "    acc = header[7]\n",
        "    f1 = header[8]\n",
        "    n_true = header[9]\n",
        "    n_pred = header[10]\n",
        "    mean_true = header[11]\n",
        "    mean_matched = header[12]\n",
        "    panoptic = header[13]\n",
        "    header = \"\"\"\n",
        "    <tr>\n",
        "    <th width = 5% align=\"center\">{0}</th>\n",
        "    <th width = 12% align=\"center\">{1}</th>\n",
        "    <th width = 6% align=\"center\">{2}</th>\n",
        "    <th width = 6% align=\"center\">{3}</th>\n",
        "    <th width = 6% align=\"center\">{4}</th>\n",
        "    <th width = 5% align=\"center\">{5}</th>\n",
        "    <th width = 5% align=\"center\">{6}</th>\n",
        "    <th width = 5% align=\"center\">{7}</th>\n",
        "    <th width = 5% align=\"center\">{8}</th>\n",
        "    <th width = 5% align=\"center\">{9}</th>\n",
        "    <th width = 5% align=\"center\">{10}</th>\n",
        "    <th width = 10% align=\"center\">{11}</th>\n",
        "    <th width = 11% align=\"center\">{12}</th>\n",
        "    <th width = 11% align=\"center\">{13}</th>\n",
        "    </tr>\"\"\".format(\"image #\",\"Prediction v. GT IoU\",'false pos.','true pos.','false neg.',precision,recall,acc,f1,n_true,n_pred,mean_true,mean_matched,panoptic)\n",
        "    html = html+header\n",
        "    i=0\n",
        "    for row in metrics:\n",
        "      i+=1\n",
        "      #image = row[0]\n",
        "      PvGT_IoU = row[1]\n",
        "      fp = row[2]\n",
        "      tp = row[3]\n",
        "      fn = row[4]\n",
        "      precision = row[5]\n",
        "      recall = row[6]\n",
        "      acc = row[7]\n",
        "      f1 = row[8]\n",
        "      n_true = row[9]\n",
        "      n_pred = row[10]\n",
        "      mean_true = row[11]\n",
        "      mean_matched = row[12]\n",
        "      panoptic = row[13]\n",
        "      cells = \"\"\"\n",
        "        <tr>\n",
        "          <td>{0}</td>\n",
        "          <td>{1}</td>\n",
        "          <td>{2}</td>\n",
        "          <td>{3}</td>\n",
        "          <td>{4}</td>\n",
        "          <td>{5}</td>\n",
        "          <td>{6}</td>\n",
        "          <td>{7}</td>\n",
        "          <td>{8}</td>\n",
        "          <td>{9}</td>\n",
        "          <td>{10}</td>\n",
        "          <td>{11}</td>\n",
        "          <td>{12}</td>\n",
        "          <td>{13}</td>\n",
        "        </tr>\"\"\".format(str(i),str(round(float(PvGT_IoU),3)),fp,tp,fn,str(round(float(precision),3)),str(round(float(recall),3)),str(round(float(acc),3)),str(round(float(f1),3)),n_true,n_pred,str(round(float(mean_true),3)),str(round(float(mean_matched),3)),str(round(float(panoptic),3)))\n",
        "      html = html+cells\n",
        "    html = html+\"\"\"</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 = '- StarDist 2D: 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",
        "\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",
        "# Function to check source and target file paths include base path\n",
        "def check_base_path(base_path, path_to_data):\n",
        "  if base_path not in path_to_data:\n",
        "    if path_to_data[0] == '/':\n",
        "      path_to_data = path_to_data[1:]\n",
        "    path_to_data = os.path.join(base_path, path_to_data)\n",
        "  return path_to_data\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",
        "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 instructions.\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",
        "\n",
        "from google.colab import drive\n",
        "drive.mount('/content/gdrive')\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xyjdhq0Zt5TF"
      },
      "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": "b_94owZCpv0v"
      },
      "source": [
        "## **3.1. Download a test dataset**\n",
        "---\n",
        "\n",
        "<font size=4>If you don't have a dataset to use with this notebook, you can download a test dataset by running the following cell.</font>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "P1f_hZ_ApuNZ"
      },
      "outputs": [],
      "source": [
        "#@markdown ##Download a test dataset \n",
        "\n",
        "\n",
        "import requests\n",
        "import zipfile\n",
        "import io\n",
        "\n",
        "# Define the Zenodo link\n",
        "zenodo_link = \"https://zenodo.org/record/7949940/files/Stardist.zip?download=1\"\n",
        "\n",
        "# Download the file\n",
        "response = requests.get(zenodo_link)\n",
        "zip_file = zipfile.ZipFile(io.BytesIO(response.content))\n",
        "\n",
        "# Specify the destination folder for extraction\n",
        "destination_folder = base_path + \"//data\"\n",
        "\n",
        "\n",
        "# Extract the contents of the zip file to the destination folder\n",
        "zip_file.extractall(destination_folder)\n",
        "\n",
        "# Close the zip file\n",
        "zip_file.close()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FQ_QxtSWQ7CL"
      },
      "source": [
        "## **3.2. 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",
        "\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",
        "\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 50-100 epochs, but a full training should run for up to 400 epochs. Evaluate the performance after training (see 5.). **Default value: 100**\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: 2**\n",
        "\n",
        "<font size = 4>**`number_of_steps`:** Define the number of training steps by epoch. By default this parameter is calculated so that each image / patch is seen at least once per epoch. **Default value: Number of patch / batch_size**\n",
        "\n",
        "<font size = 4>**`patch_size`:** Input the size of the patches use to train StarDist 2D (length of a side). The value should be smaller or equal to the dimensions of the image. Make the patch size as large as possible and divisible by 8. **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 square has 4 corners). **Default value: 32** \n",
        "\n",
        "<font size = 4>**`grid_parameter`:** increase this number if the cells/nuclei are very large or decrease it if they are very small. **Default value: 2**\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 value until the OOM error disappear.**\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "ewpNJ_I0Mv47"
      },
      "outputs": [],
      "source": [
        "#@markdown ###Path to training images: \n",
        "Training_source = '/data/Stardist/Training - Images' #@param {type:\"string\"}\n",
        "\n",
        "Training_target = '/data/Stardist/Training - Masks' #@param {type:\"string\"}\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",
        "\n",
        "#@markdown ### Other parameters for training:\n",
        "number_of_epochs =  50#@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 =  15#@param {type:\"number\"}\n",
        "number_of_steps = 50#@param {type:\"number\"}\n",
        "patch_size = 512 #@param {type:\"number\"}\n",
        "percentage_validation =  10 #@param {type:\"number\"}\n",
        "n_rays = 32 #@param {type:\"number\"}\n",
        "grid_parameter = 2#@param [1, 2, 4, 8, 16, 32] {type:\"raw\"}\n",
        "initial_learning_rate = 0.0001 #@param {type:\"number\"}\n",
        "\n",
        "if (Use_Default_Advanced_Parameters): \n",
        "  print(\"Default advanced parameters enabled\")\n",
        "  batch_size = 2\n",
        "  n_rays = 32\n",
        "  percentage_validation = 10\n",
        "  grid_parameter = 2\n",
        "  initial_learning_rate = 0.0003\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 !!\")\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",
        "# Check if the Training_source and Training_target folders contain the base_path variable, if not add it\n",
        "Training_source = check_base_path(base_path, Training_source)\n",
        "Training_target = check_base_path(base_path, Training_target)\n",
        "\n",
        "# Here we open will randomly chosen input and output image\n",
        "random_choice = random.choice(os.listdir(Training_source))\n",
        "x = imread(Training_source+\"/\"+random_choice)\n",
        "\n",
        "# Here we check the image dimensions\n",
        "\n",
        "Image_Y = x.shape[0]\n",
        "Image_X = x.shape[1]\n",
        "\n",
        "print('Loaded images (width, length) =', x.shape)\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",
        "  \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",
        "if patch_size > 2048:\n",
        "  patch_size = 2048\n",
        "  print(bcolors.WARNING + \" Your image dimension is large; therefore the patch_size chosen is now:\",patch_size)\n",
        "\n",
        "\n",
        "# Here we check that the patch_size is divisible by 16\n",
        "if not patch_size % 16 == 0:\n",
        "    patch_size = ((int(patch_size / 16)-1) * 16)\n",
        "    print(bcolors.WARNING + \" Your chosen patch_size is not divisible by 8; therefore the patch_size chosen is:\",patch_size)\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",
        "\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",
        "norm = simple_norm(x, percent = 99)\n",
        "\n",
        "f=plt.figure(figsize=(16,8))\n",
        "plt.subplot(1,2,1)\n",
        "plt.imshow(x, interpolation='nearest', norm=norm, cmap='magma')\n",
        "plt.title('Training source')\n",
        "plt.axis('off');\n",
        "\n",
        "plt.subplot(1,2,2)\n",
        "plt.imshow(y, interpolation='nearest', cmap=lbl_cmap)\n",
        "plt.title('Training target')\n",
        "plt.axis('off');\n",
        "plt.savefig(base_path + '/TrainingDataExample_StarDist2D.png',bbox_inches='tight',pad_inches=0)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xyQZKby8yFME"
      },
      "source": [
        "## **3.3. 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>Data augmentation is performed here via random rotations, flips, and intensity changes.\n",
        "\n",
        "\n",
        "<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.** "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "DMqWq5-AxnFU"
      },
      "outputs": [],
      "source": [
        "#Data augmentation\n",
        "\n",
        "Use_Data_augmentation = False #@param {type:\"boolean\"}\n",
        "\n",
        "#@markdown ####Choose a factor by which you want to multiply your original dataset\n",
        "\n",
        "Multiply_dataset_by = 4 #@param {type:\"slider\", min:1, max:10, step:1}\n",
        "\n",
        "\n",
        "def random_fliprot(img, mask): \n",
        "    assert img.ndim >= mask.ndim\n",
        "    axes = tuple(range(mask.ndim))\n",
        "    perm = tuple(np.random.permutation(axes))\n",
        "    img = img.transpose(perm + tuple(range(mask.ndim, img.ndim))) \n",
        "    mask = mask.transpose(perm) \n",
        "    for ax in axes: \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",
        "\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",
        "    x, y = random_fliprot(x, y)\n",
        "    x = random_intensity_change(x)\n",
        "    # add some gaussian noise\n",
        "    sig = 0.02*np.random.uniform(0,1)\n",
        "    x = x + sig*np.random.normal(0,1,x.shape)\n",
        "    return x, y\n",
        "\n",
        "\n",
        "\n",
        "if Use_Data_augmentation:\n",
        "  augmenter = augmenter\n",
        "  print(\"Data augmentation enabled\")\n",
        "\n",
        "\n",
        "if not Use_Data_augmentation:\n",
        "  augmenter = None\n",
        "  print(bcolors.WARNING+\"Data augmentation disabled\") \n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3L9zSGtORKYI"
      },
      "source": [
        "\n",
        "## **3.4. Using weights from a pre-trained model as initial weights**\n",
        "---\n",
        "<font size = 4>  Here, you can set the path to a pre-trained model or the ID of a model from the [BioImage Model Zoo](https://bioimage.io/) (*e.g.*, 10.5281/zenodo.5749843) 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",
        "pretrained_model_choice = \"2D_versatile_fluo_from_Stardist_Fiji\" #@param [\"Model_from_file\", \"BioImage Model Zoo\", \"2D_versatile_fluo_from_Stardist_Fiji\", \"2D_Demo_Model_from_Stardist_Github\", \"Versatile_H&E_nuclei\"]\n",
        "Weights_choice = \"best\" #@param [\"last\", \"best\"]\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",
        "#@markdown ###If you chose \"BioImage Model Zoo\", please provide the path or the URL to the model:\n",
        "bioimageio_model = \"\" #@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",
        "  # --------------------- Load the model from the choosen path ------------------------\n",
        "  elif pretrained_model_choice == \"BioImage Model Zoo\":\n",
        "    h5_file_path = base_path + \"/bioimage_model_zoo/\"\n",
        "    %rm -rf $h5_file_path\n",
        "    model_imported = import_bioimageio(bioimageio_model, h5_file_path)\n",
        "    del model_imported\n",
        "    biomodel = load_resource_description(os.path.join(h5_file_path, \"bioimageio\", \"rdf.yaml\"))   \n",
        "    h5_file_path = os.path.join(h5_file_path,  \"bioimageio\", biomodel.config['stardist']['weights'])\n",
        "\n",
        "# --------------------- Download the Demo 2D model provided in the Stardist 2D github ------------------------\n",
        "\n",
        "  if pretrained_model_choice == \"2D_Demo_Model_from_Stardist_Github\":\n",
        "    pretrained_model_name = \"2D_Demo\"\n",
        "    pretrained_model_path = base_path + \"/\" + pretrained_model_name\n",
        "    print(\"Downloading the 2D_Demo_Model_from_Stardist_Github\")\n",
        "    if os.path.exists(pretrained_model_path):\n",
        "      shutil.rmtree(pretrained_model_path)\n",
        "    os.makedirs(pretrained_model_path)\n",
        "    wget.download(\"https://github.com/mpicbg-csbd/stardist/raw/master/models/examples/2D_demo/config.json\", pretrained_model_path)\n",
        "    wget.download(\"https://github.com/mpicbg-csbd/stardist/raw/master/models/examples/2D_demo/thresholds.json\", pretrained_model_path)\n",
        "    wget.download(\"https://github.com/mpicbg-csbd/stardist/blob/master/models/examples/2D_demo/weights_best.h5?raw=true\", pretrained_model_path)    \n",
        "    wget.download(\"https://github.com/mpicbg-csbd/stardist/blob/master/models/examples/2D_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",
        "# --------------------- Download the Demo 2D_versatile_fluo_from_Stardist_Fiji ------------------------\n",
        "\n",
        "  if pretrained_model_choice == \"2D_versatile_fluo_from_Stardist_Fiji\":\n",
        "    print(\"Downloading the 2D_versatile_fluo_from_Stardist_Fiji\")\n",
        "    pretrained_model_name = \"2D_versatile_fluo\"\n",
        "    pretrained_model_path = base_path + \"/\" + pretrained_model_name\n",
        "    \n",
        "    if os.path.exists(pretrained_model_path):\n",
        "      shutil.rmtree(pretrained_model_path)\n",
        "    os.makedirs(pretrained_model_path)\n",
        "    \n",
        "    wget.download(\"https://cloud.mpi-cbg.de/index.php/s/1k5Zcy7PpFWRb0Q/download?path=/versatile&files=2D_versatile_fluo.zip\", pretrained_model_path)\n",
        "    \n",
        "    with zipfile.ZipFile(pretrained_model_path+\"/2D_versatile_fluo.zip\", 'r') as zip_ref:\n",
        "      zip_ref.extractall(pretrained_model_path)\n",
        "    \n",
        "    h5_file_path = os.path.join(pretrained_model_path, \"weights_best.h5\")\n",
        "\n",
        "# --------------------- Download the Versatile (H&E nuclei)_fluo_from_Stardist_Fiji ------------------------\n",
        "\n",
        "  if pretrained_model_choice == \"Versatile_H&E_nuclei\":\n",
        "    print(\"Downloading the Versatile_H&E_nuclei from_Stardist_Fiji\")\n",
        "    pretrained_model_name = \"2D_versatile_he\"\n",
        "    pretrained_model_path = base_path + \"/\" + pretrained_model_name\n",
        "    \n",
        "    if os.path.exists(pretrained_model_path):\n",
        "      shutil.rmtree(pretrained_model_path)\n",
        "    os.makedirs(pretrained_model_path)\n",
        "    \n",
        "    wget.download(\"https://cloud.mpi-cbg.de/index.php/s/1k5Zcy7PpFWRb0Q/download?path=/versatile&files=2D_versatile_he.zip\", pretrained_model_path)\n",
        "    \n",
        "    with zipfile.ZipFile(pretrained_model_path+\"/2D_versatile_he.zip\", 'r') as zip_ref:\n",
        "      zip_ref.extractall(pretrained_model_path)\n",
        "    \n",
        "    h5_file_path = os.path.join(pretrained_model_path, \"weights_best.h5\")\n",
        "\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('Weights found in:')\n",
        "  print(h5_file_path)\n",
        "  print('will be loaded prior to training.')\n",
        "\n",
        "else:\n",
        "  print(bcolors.WARNING+'No pretrained network will be used.')\n",
        "\n"
      ]
    },
    {
      "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",
        "# --------------------- Here we load the augmented data or the raw data ------------------------\n",
        "\n",
        "\n",
        "Training_source_dir = Training_source\n",
        "Training_target_dir = Training_target\n",
        "# --------------------- ------------------------------------------------\n",
        "\n",
        "training_images_tiff=Training_source_dir+\"/*.tif\"\n",
        "mask_images_tiff=Training_target_dir+\"/*.tif\"\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_channel = 1 if X[0].ndim == 2 else X[0].shape[-1]\n",
        "\n",
        "# --------------------- Piece of code to remove alpha channel in RGBA images ------------------------\n",
        "\n",
        "if n_channel == 4:\n",
        "  X[:] = [i[:,:,:3] for i in X]\n",
        "  print(\"The alpha channel has been removed\")\n",
        "  n_channel = 1 if X[0].ndim == 2 else X[0].shape[-1]\n",
        "\n",
        "\n",
        "if not Use_Data_augmentation:\n",
        "  augmenter = None\n",
        "  \n",
        "\n",
        "#Normalize images and fill small label holes.\n",
        "axis_norm = (0,1) \n",
        "if n_channel == 1:\n",
        "  # normalize channels independently\n",
        "  print(\"Normalizing image channels independently\")\n",
        "if n_channel > 1:\n",
        "  # normalize channels independently\n",
        "  # axis_norm = (0,1,2) # normalize channels jointly\n",
        "  print(\"Normalizing image channels %s.\" % ('jointly' if axis_norm is None or 2 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",
        "#It is advisable to use 10 % of your training dataset for validation. This ensures the truthfull validation error value. If only few validation images are used network may choose too easy or too challenging images for validation. \n",
        "# split training data (images and masks) into training images and validation images.\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",
        "# Use OpenCL-based computations for data generator during training (requires 'gputools')\n",
        "# Currently always false for stability\n",
        "use_gpu = False and gputools_available()\n",
        "\n",
        "#Here we ensure that our network has a minimal number of steps\n",
        "\n",
        "if (Use_Default_Advanced_Parameters) or (number_of_steps == 0): \n",
        "  # number_of_steps= (int(len(X)/batch_size)+1)\n",
        "  number_of_steps = int(Image_X*Image_Y/(patch_size*patch_size))*(int(len(X)/batch_size)+1)\n",
        "  if (Use_Data_augmentation):\n",
        "    augmentation_factor = Multiply_dataset_by\n",
        "    number_of_steps = number_of_steps * augmentation_factor\n",
        "\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",
        "\n",
        "\n",
        "conf = Config2D (\n",
        "    n_rays       = n_rays,\n",
        "    use_gpu      = use_gpu,\n",
        "    train_batch_size = batch_size,\n",
        "    n_channel_in = n_channel,\n",
        "    train_patch_size = (patch_size, patch_size),\n",
        "    grid = (grid_parameter, grid_parameter),\n",
        "    train_learning_rate = initial_learning_rate,\n",
        ")\n",
        "\n",
        "# Here we create a model according to section 5.3.\n",
        "model = StarDist2D(conf, name=model_name, basedir=model_path)\n",
        "\n",
        "# --------------------- Using pretrained model ------------------------\n",
        "# Load the pretrained weights \n",
        "if Use_pretrained_model:\n",
        "  try:\n",
        "      model.load_weights(h5_file_path)\n",
        "  except:\n",
        "      print(bcolors.WARNING + \"The pretrained model could not be loaded as the configuration of the network is different.\")\n",
        "      print(\"Please, read the model specifications and check the parameters in Section 3.1\" + W)\n",
        "# --------------------- ---------------------- ------------------------\n",
        "\n",
        "#Here we check the FOV of the network.\n",
        "median_size = calculate_extents(list(Y), np.median)\n",
        "fov = np.array(model._axes_tile_overlap('YX'))\n",
        "if any(median_size > fov):\n",
        "    print(bcolors.WARNING+\"WARNING: median object size larger than field of view of the neural network.\")\n",
        "print(conf)\n",
        "\n",
        "pdf_export(augmentation = Use_Data_augmentation, pretrained_model = Use_pretrained_model)\n",
        "\n",
        "print(\"Number of steps: \"+str(number_of_steps))\n",
        "\n"
      ]
    },
    {
      "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": [
        "#@markdown ##Start training\n",
        "\n",
        "start = time.time()\n",
        "\n",
        "\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",
        "print(\"Network optimization in progress\")\n",
        "#Here we optimize the network.\n",
        "model.optimize_thresholds(X_val, Y_val)\n",
        "\n",
        "print(\"Done\")\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",
        "\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",
        "model.export_TF()\n",
        "\n",
        "#print(\"Your model has been sucessfully exported and can now also be used in the Stardist Fiji plugin\")\n",
        "\n",
        "pdf_export(trained=True, augmentation = Use_Data_augmentation, pretrained_model = Use_pretrained_model)\n",
        "\n",
        "#Create a pdf document with training summary"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_0Hynw3-xHp1"
      },
      "source": [
        "# **5. Evaluate your model**\n",
        "---\n",
        "\n",
        "<font size = 4>This section allows you to perform important quality checks on the validity and generalisability of the trained model.  \n",
        "\n",
        "\n",
        "<font size = 4>**We highly recommend to perform quality control on all newly trained models.**\n",
        "\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",
        "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",
        "  print(bcolors.WARNING+'!! WARNING: The chosen model does not exist !!')\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.\n",
        "\n",
        "\n"
      ]
    },
    {
      "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",
        "\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",
        "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** (IuO) 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>Here, the IuO is both calculated over the whole image and on a per-object basis. The value displayed below is the IuO value calculated over the entire image. The IuO value calculated on a per-object basis is used to calculate the other metrics displayed.\n",
        "\n",
        "<font size = 4>“n_true” refers to the number of objects present in the ground truth image. “n_pred” refers to the number of objects present in the predicted image. \n",
        "\n",
        "<font size = 4>When a segmented object has an IuO value above 0.5 (compared to the corresponding ground truth), it is then considered a true positive. The number of “**true positives**” is available in the table below. The number of “false positive” is then defined as  “**false positive**” = “n_pred” - “true positive”. The number of “false negative” is defined as “false negative” = “n_true” - “true positive”.\n",
        "\n",
        "<font size = 4>The mean_matched_score is the mean IoUs of matched true positives. The mean_true_score is the mean IoUs of matched true positives but normalized by the total number of ground truth objects. The panoptic_quality is calculated as described by [Kirillov et al. 2019](https://arxiv.org/abs/1801.00868).\n",
        "\n",
        "<font size = 4>For more information about the other metric displayed, please consult the SI of the paper describing ZeroCostDL4Mic.\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 ##Choose the folders that contain your Quality Control dataset\n",
        "\n",
        "from stardist.matching import matching\n",
        "from stardist.plot import render_label, render_label_pred \n",
        "\n",
        "Source_QC_folder = '/data/Stardist/Test - Images' #@param{type:\"string\"}\n",
        "Target_QC_folder = '/data/Stardist/Test - Masks' #@param{type:\"string\"}\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",
        "# Check if the sourcce and target folders contain base_path, if not add it\n",
        "Source_QC_folder = check_base_path(base_path, Source_QC_folder)\n",
        "Target_QC_folder = check_base_path(base_path, Target_QC_folder)\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",
        "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",
        "\n",
        "print('Number of test dataset found in the folder: '+str(len(Z)))\n",
        "  \n",
        "#Normalize images.\n",
        "\n",
        "if n_channel == 1:\n",
        "  axis_norm = (0,1)   # normalize channels independently\n",
        "  print(\"Normalizing image channels independently\")\n",
        "\n",
        "if n_channel > 1:\n",
        "  axis_norm = (0,1,2) # normalize channels jointly\n",
        "  print(\"Normalizing image channels jointly\") \n",
        "\n",
        "model = StarDist2D(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",
        "  os.chdir(QC_model_path+\"/\"+QC_model_name+\"/Quality Control/Prediction\")\n",
        "  imwrite(names[i], labels)\n",
        "\n",
        "# Here we start testing the differences between GT and predicted masks\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, delimiter=\",\")\n",
        "  writer.writerow([\"image\",\"Prediction v. GT Intersection over Union\", \"false positive\", \"true positive\", \"false negative\", \"precision\", \"recall\", \"accuracy\", \"f1 score\", \"n_true\", \"n_pred\", \"mean_true_score\", \"mean_matched_score\", \"panoptic_quality\"])  \n",
        "\n",
        "# define the images\n",
        "\n",
        "  for n in os.listdir(Source_QC_folder):\n",
        "    \n",
        "    if not os.path.isdir(os.path.join(Source_QC_folder,n)):\n",
        "      print('Running QC on: '+n)\n",
        "      test_input = imread(os.path.join(Source_QC_folder,n))\n",
        "      test_prediction = imread(os.path.join(QC_model_path+\"/\"+QC_model_name+\"/Quality Control/Prediction\",n))\n",
        "      test_ground_truth_image = imread(os.path.join(Target_QC_folder, n))\n",
        "\n",
        "      # Calculate the matching (with IoU threshold `thresh`) and all metrics\n",
        "      \n",
        "      stats = matching(test_ground_truth_image, test_prediction, thresh=0.5)\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",
        "\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), str(stats.fp), str(stats.tp), str(stats.fn), str(stats.precision), str(stats.recall), str(stats.accuracy), str(stats.f1), str(stats.n_true), str(stats.n_pred), str(stats.mean_true_score), str(stats.mean_matched_score), str(stats.panoptic_quality)])\n",
        "\n",
        "from tabulate import tabulate\n",
        "\n",
        "df = pd.read_csv (QC_model_path+\"/\"+QC_model_name+\"/Quality Control/Quality_Control for \"+QC_model_name+\".csv\")\n",
        "print(tabulate(df, headers='keys', tablefmt='psql'))\n",
        "\n",
        "\n",
        "from astropy.visualization import simple_norm\n",
        "\n",
        "# ------------- For display ------------\n",
        "print('--------------------------------------------------------------')\n",
        "@interact\n",
        "def show_QC_results(file = os.listdir(Source_QC_folder)):\n",
        "  \n",
        "\n",
        "  plt.figure(figsize=(25,5))\n",
        "  if n_channel > 1:\n",
        "    source_image = imread(os.path.join(Source_QC_folder, file))\n",
        "  if n_channel == 1:\n",
        "    source_image = imread(os.path.join(Source_QC_folder, file))\n",
        "\n",
        "  target_image = imread(os.path.join(Target_QC_folder, file))\n",
        "  prediction = imread(QC_model_path+\"/\"+QC_model_name+\"/Quality Control/Prediction/\"+file)\n",
        "\n",
        "  stats = matching(prediction, target_image, thresh=0.5)\n",
        "\n",
        "  target_image_mask = np.empty_like(target_image)\n",
        "  target_image_mask[target_image > 0] = 255\n",
        "  target_image_mask[target_image == 0] = 0\n",
        "  \n",
        "  prediction_mask = np.empty_like(prediction)\n",
        "  prediction_mask[prediction > 0] = 255\n",
        "  prediction_mask[prediction == 0] = 0\n",
        "\n",
        "  intersection = np.logical_and(target_image_mask, prediction_mask)\n",
        "  union = np.logical_or(target_image_mask, prediction_mask)\n",
        "  iou_score =  np.sum(intersection) / np.sum(union)\n",
        "\n",
        "  norm = simple_norm(source_image, percent = 99)\n",
        "\n",
        "  #Input\n",
        "  plt.subplot(1,4,1)\n",
        "  plt.axis('off')\n",
        "  if n_channel > 1:\n",
        "    plt.imshow(source_image)\n",
        "  if n_channel == 1:\n",
        "    plt.imshow(source_image, 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(target_image_mask, 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(prediction_mask, aspect='equal', cmap='Purples')\n",
        "  plt.title('Prediction')\n",
        "\n",
        "  #Overlay\n",
        "  plt.subplot(1,4,4)\n",
        "  plt.axis('off')\n",
        "  plt.imshow(target_image_mask, cmap='Greens')\n",
        "  plt.imshow(prediction_mask, alpha=0.5, cmap='Purples')\n",
        "  plt.title('Ground Truth and Prediction, Intersection over Union:'+str(round(iou_score,3 )));\n",
        "  plt.savefig(full_QC_model_path+'/Quality Control/QC_example_data.png',bbox_inches='tight',pad_inches=0)\n",
        "\n",
        "\n",
        "qc_pdf_export()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7j1VuFCs_g5P"
      },
      "source": [
        "## **5.3. Export your model into the BioImage Model Zoo format**\n",
        "---\n",
        "<font size = 4>This section exports the model into the [BioImage Model Zoo](https://bioimage.io/#/) format so it can be used directly with deepImageJ or Ilastik. The new files will be stored in the model folder specified at the beginning of Section 5. \n",
        "\n",
        "<font size = 4>Once the cell is executed, you will find a new zip file with the name specified in `trained_model_name.bioimage.io.model`.\n",
        "\n",
        "<font size = 4>To use it with deepImageJ, download it and install it suing DeepImageJ Install Model > Install from a local file. \n",
        "\n",
        "<font size = 4>To try the model in ImageJ, go to Plugins > DeepImageJ > DeepImageJ Run, choose this model from the list and click on Test Model.\n",
        "\n",
        "<font size = 4>The exported model contains an additional ImageJ macro (`StarDist2D_Post-processing.ijm`) to run the StarDist postprocessing in Fiji.\n",
        "\n",
        "<font size = 4>  More information at https://deepimagej.github.io/deepimagej/"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "v5F2VWnQ_fL-"
      },
      "outputs": [],
      "source": [
        "# ------------- User input ------------\n",
        "os.chdir(base_path + \"/\")\n",
        "\n",
        "# information about the model\n",
        "#@markdown ##Introduce the metadata of the model architecture:\n",
        "Trained_model_name    = \"\" #@param {type:\"string\"}\n",
        "Trained_model_authors =  \"[Author 1, Author 2, Author 3]\" #@param {type:\"string\"}\n",
        "Trained_model_authors_affiliation =  \"[Affiliation 1, Affiliation 2, Affiliation 3]\" #@param {type:\"string\"}\n",
        "Trained_model_description = \"\" #@param {type:\"string\"}\n",
        "Trained_model_license = 'MIT'#@param {type:\"string\"}\n",
        "Trained_model_references = [ \"Uwe Schmidt et al. MICCAI 2018\", \"Martin Weigert et al. WACV 2020\", \"Lucas von Chamier et al. biorXiv 2020\"] \n",
        "Trained_model_DOI = [\"https://doi.org/10.1007/978-3-030-00934-2_30\", \"https://doi.org/10.1109/WACV45572.2020.9093435\", \"https://doi.org/10.1101/2020.03.20.000133\"] \n",
        "\n",
        "#@markdown ##Introduce the pixel size (in microns) of the image provided as an example of the model processing:\n",
        "# information about the example image\n",
        "PixelSize = 1 #@param {type:\"number\"}\n",
        "#@markdown ##Do you want to choose the exampleimage?\n",
        "default_example_image = True #@param {type:\"boolean\"}\n",
        "#@markdown ###If not, please input:\n",
        "fileID    =  \"\" #@param {type:\"string\"}\n",
        "\n",
        "\n",
        "# Load an example image\n",
        "if default_example_image:\n",
        "    i=0 # first image of the target list\n",
        "    test_img = Z[i]\n",
        "else:\n",
        "    test_img = imread(fileID)\n",
        "\n",
        "# Load the model\n",
        "model = StarDist2D(None, name=QC_model_name, basedir=QC_model_path)\n",
        "modelfile = Path(os.path.join(QC_model_path, QC_model_name + 'bioimageio.zip'))\n",
        "\n",
        "## The following functions should be provisional until StarDist+CSBDeep migrate\n",
        "# completely to TF2\n",
        "def _is_power_of_2(i):\n",
        "    assert i > 0\n",
        "    e = np.log2(i)\n",
        "    return e == int(e)\n",
        "\n",
        "def _my_compute_receptive_field(model, img_size=None):\n",
        "    # TODO: good enough?\n",
        "    from scipy.ndimage import zoom\n",
        "    if img_size is None:\n",
        "        img_size = tuple(g*(128 if model.config.n_dim==2 else 64) for g in model.config.grid)\n",
        "    if np.isscalar(img_size):\n",
        "        img_size = (img_size,) * model.config.n_dim\n",
        "    img_size = tuple(img_size)\n",
        "    # print(img_size)\n",
        "    assert all(_is_power_of_2(s) for s in img_size)\n",
        "    mid = tuple(s//2 for s in img_size)\n",
        "    x = np.zeros((1,)+img_size+(model.config.n_channel_in,), dtype=np.float32)\n",
        "    z = np.zeros_like(x)\n",
        "    x[(0,)+mid+(slice(None),)] = 1\n",
        "    #y  = self.keras_model.predict(x)[0][0,...,0]\n",
        "    #y0 = self.keras_model.predict(z)[0][0,...,0]\n",
        "    y  = model.predict(np.squeeze(x))[0]\n",
        "    y0 = model.predict(np.squeeze(z))[0]\n",
        "\n",
        "    grid = tuple((np.array(x.shape[1:-1])/np.array(y.shape)).astype(int))\n",
        "    assert grid == model.config.grid\n",
        "    y  = zoom(y, grid,order=0)\n",
        "    y0 = zoom(y0,grid,order=0)\n",
        "    ind = np.where(np.abs(y-y0)>0)\n",
        "    return [(m-np.min(i), np.max(i)-m) for (m,i) in zip(mid,ind)]\n",
        "\n",
        "model._tile_overlap = _my_compute_receptive_field(model)\n",
        "\n",
        "###\n",
        "export_bioimageio(model, modelfile, test_img)"
      ]
    },
    {
      "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": [
        "\n",
        "\n",
        "## **6.1 Generate prediction(s) from unseen dataset**\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.\n",
        "\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 will train.\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",
        "<font size = 4>In stardist the following results can be exported:\n",
        "- Region of interest (ROI) that can be opened in ImageJ / Fiji. The ROI are saved inside of a .zip file in your choosen result folder. To open the ROI in Fiji, just drag and drop the zip file !**\n",
        "- The predicted mask images\n",
        "- A CSV file that contains the coordinate the centre of each detected nuclei (single image only). \n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "y2TD5p7MZrEb"
      },
      "outputs": [],
      "source": [
        "Single_Images = 1\n",
        "Stacks = 2\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",
        "Results_folder = \"\" #@param {type:\"string\"}\n",
        "\n",
        "#@markdown ###Are your data single images or stacks?\n",
        "\n",
        "Data_type = Single_Images #@param [\"Single_Images\", \"Stacks\"] {type:\"raw\"}\n",
        "\n",
        "#@markdown ###Do you also want to export Fiji ROI?\n",
        "Region_of_interests = False #@param {type:\"boolean\"}\n",
        "\n",
        "# model name and path\n",
        "#@markdown ###Do you want to use the current trained model?\n",
        "Use_the_current_trained_model = False #@param {type:\"boolean\"}\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",
        "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",
        "  print(bcolors.WARNING+'!! 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",
        "#single images\n",
        "\n",
        "if Data_type == 1 :\n",
        "\n",
        "  Data_folder = Data_folder+\"/*.tif\"\n",
        "\n",
        "  print(\"Single images are now beeing predicted\")\n",
        "  np.random.seed(16)\n",
        "  lbl_cmap = random_label_cmap()\n",
        "  X = sorted(glob(Data_folder))\n",
        "  X = list(map(imread,X))\n",
        "  n_channel = 1 if X[0].ndim == 2 else X[0].shape[-1]\n",
        "  \n",
        "  # axis_norm = (0,1,2) # normalize channels jointly\n",
        "  if n_channel == 1:\n",
        "    axis_norm = (0,1)   # normalize channels independently\n",
        "    print(\"Normalizing image channels independently\")\n",
        "\n",
        "  if n_channel > 1:\n",
        "    axis_norm = (0,1,2) # normalize channels jointly\n",
        "    print(\"Normalizing image channels jointly\")  \n",
        "    sys.stdout.flush()  \n",
        "  \n",
        "  model = StarDist2D(None, name = Prediction_model_name, basedir = Prediction_model_path)\n",
        " \n",
        "  names = [os.path.basename(f) for f in sorted(glob(Data_folder))]  \n",
        "  Nuclei_number = []\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",
        "  # Create a list of name with no extension\n",
        " \n",
        "  name_no_extension=[]\n",
        "  for n in names:\n",
        "    name_no_extension.append(os.path.splitext(n)[0])\n",
        "    \n",
        "  # Save all ROIs and masks into results folder\n",
        "  \n",
        "  for i in range(len(X)):\n",
        "    \n",
        "    img = normalize(X[i], 1,99.8, axis = axis_norm)\n",
        "    labels, polygons = model.predict_instances(img)\n",
        "    \n",
        "    os.chdir(Results_folder)\n",
        "    imwrite(FILEnames[i], labels)\n",
        "    export_imagej_rois(name_no_extension[i], polygons['coord'])      \n",
        "    Nuclei_centre_coordinate = polygons['points']\n",
        "    my_df2 = pd.DataFrame(Nuclei_centre_coordinate)\n",
        "    my_df2.columns =['Y', 'X']\n",
        "    \n",
        "    my_df2.to_csv(Results_folder+'/'+name_no_extension[i]+'_Nuclei_centre.csv', index=False, header=True)\n",
        "\n",
        "    Nuclei_array = polygons['coord']\n",
        "    Nuclei_array2 = [names[i], Nuclei_array.shape[0]]\n",
        "    Nuclei_number.append(Nuclei_array2) \n",
        "\n",
        "  my_df = pd.DataFrame(Nuclei_number)\n",
        "  my_df2.columns =['Frame number', 'Number of objects']\n",
        "  my_df.to_csv(Results_folder+'/Nuclei_count.csv', index=False, header=False)\n",
        "  \n",
        "# One example is displayed\n",
        "\n",
        "  print(\"One example image is displayed bellow:\")\n",
        "  plt.figure(figsize=(10,10))\n",
        "  plt.imshow(img if img.ndim==2 else img[...,:3], clim=(0,1), cmap='gray')\n",
        "  plt.imshow(labels, cmap=lbl_cmap, alpha=0.5)\n",
        "  plt.axis('off');\n",
        "  plt.savefig(name_no_extension[i]+\"_overlay.tif\")\n",
        "\n",
        "# Here is the code to analyse stacks\n",
        "\n",
        "if Data_type == 2 :\n",
        "  print(\"Stacks are now beeing predicted\")\n",
        "  np.random.seed(42)\n",
        "  lbl_cmap = random_label_cmap()\n",
        "\n",
        "  # normalize channels independently\n",
        "  axis_norm = (0,1)  \n",
        "  \n",
        "  model = StarDist2D(None, name = Prediction_model_name, basedir = Prediction_model_path)\n",
        "  \n",
        "  for image in os.listdir(Data_folder):\n",
        "    print(\"Performing prediction on: \"+image)\n",
        "\n",
        "    Number_of_nuclei_list = []\n",
        "    Number_of_frame_list = []\n",
        "\n",
        "    timelapse = imread(Data_folder+\"/\"+image)\n",
        "\n",
        "    short_name = os.path.splitext(image)     \n",
        "    \n",
        "    timelapse = normalize(timelapse, 1,99.8, axis=(0,)+tuple(1+np.array(axis_norm)))\n",
        "    \n",
        "\n",
        "    if Region_of_interests: \n",
        "      polygons = [model.predict_instances(frame)[1]['coord'] for frame in tqdm(timelapse)]    \n",
        "      export_imagej_rois(Results_folder+\"/\"+str(short_name[0]), polygons, compression=ZIP_DEFLATED)  \n",
        "\n",
        "\n",
        "# Analyse each time points one after the other\n",
        "    \n",
        "    n_timepoint = timelapse.shape[0]\n",
        "    prediction_stack = np.zeros((n_timepoint, timelapse.shape[1], timelapse.shape[2]))\n",
        "    \n",
        "    for t in range(n_timepoint):\n",
        "      img_t = timelapse[t]\n",
        "      labels, polygons = model.predict_instances(img_t)      \n",
        "      prediction_stack[t] = labels\n",
        "      Nuclei_array = polygons['coord']\n",
        "      Nuclei_array2 = [str(t), Nuclei_array.shape[0]]\n",
        "      Number_of_nuclei_list.append(Nuclei_array2)\n",
        "      Number_of_frame_list.append(t)\n",
        "\n",
        "    #prediction_stack_32 = img_as_float32(prediction_stack, force_copy=False)\n",
        "    prediction_stack = img_as_float32(prediction_stack, force_copy=False)\n",
        "\n",
        "# Export a csv file containing the number of nuclei detected at each frame\n",
        "    my_df = pd.DataFrame(Number_of_nuclei_list)\n",
        "    my_df.to_csv(Results_folder+'/'+str(short_name[0])+'_Nuclei_number.csv', index=False, header=False)\n",
        "    os.chdir(Results_folder)      \n",
        "    imsave(str(short_name[0])+\".tif\", prediction_stack, compression ='zlib')\n",
        "\n",
        "  # Object detected vs frame number\n",
        "    plt.figure(figsize=(20,5))\n",
        "    my_df.plot()\n",
        "    plt.title('Number of objects vs frame number')\n",
        "    plt.ylabel('Number of detected objects')\n",
        "    plt.xlabel('Frame number')\n",
        "    plt.legend()\n",
        "    plt.savefig(Results_folder+'/'+str(short_name[0])+'_Object_detected_vs_frame_number.png',bbox_inches='tight',pad_inches=0)\n",
        "    plt.show()\n",
        "    del prediction_stack\n",
        "    del timelapse\n",
        "     \n",
        "\n",
        "print(\"Predictions completed\")   "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KjGHBGmxlk9B"
      },
      "source": [
        "## **6.2. Generate prediction(s) from unseen dataset (Big data)**\n",
        "---\n",
        "\n",
        "<font size = 4>You can use this section of the notebook to generate predictions on very large images. Compatible file formats include .Tif and .svs files.\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",
        "\n",
        "<font size = 4>In stardist the following results can be exported:\n",
        "- Region of interest (ROI) that can be opened in ImageJ / Fiji. The ROI are saved inside of a .zip file in your choosen result folder. To open the ROI in Fiji, just drag and drop the zip file ! IMPORTANT: ROI files cannot be exported for extremely large images.\n",
        "- The predicted mask images\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "jxjHeOFFleSV"
      },
      "outputs": [],
      "source": [
        "#@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",
        "\n",
        "start = time.time()\n",
        "\n",
        "\n",
        "Data_folder = \"\" #@param {type:\"string\"}\n",
        "\n",
        "Results_folder = \"\" #@param {type:\"string\"}\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",
        "\n",
        "#@markdown #####To analyse very large image, your images need to be divided into blocks.  Each blocks will then be processed independently and re-assembled to generate the final image. \n",
        "#@markdown #####Here you can choose the dimension of the block.\n",
        "\n",
        "block_size_Y =   1024#@param {type:\"number\"}\n",
        "block_size_X =   1024#@param {type:\"number\"}\n",
        "\n",
        "\n",
        "#@markdown #####Here you can the amount of overlap between each block.\n",
        "min_overlap =   50#@param {type:\"number\"}\n",
        "\n",
        "#@markdown #####To analyse large blocks, your blocks need to be divided into tiles.  Each tile will then be processed independently and re-assembled to generate the final block. \n",
        "\n",
        "n_tiles_Y =  1#@param {type:\"number\"}\n",
        "n_tiles_X =  1#@param {type:\"number\"}\n",
        "\n",
        "\n",
        "#@markdown ###What outputs would you like to generate? \n",
        "Mask_images = True #@param {type:\"boolean\"}\n",
        "\n",
        "Region_of_interests = True #@param {type:\"boolean\"}\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",
        "\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",
        "#Create a temp folder to save Zarr files\n",
        "\n",
        "Temp_folder = base_path + \"/Temp_folder\"\n",
        "\n",
        "if os.path.exists(Temp_folder):\n",
        "  shutil.rmtree(Temp_folder)\n",
        "os.makedirs(Temp_folder)\n",
        "\n",
        "\n",
        "# mi, ma = np.percentile(img[::8], [1,99.8])                      # compute percentiles from low-resolution image\n",
        "# mi, ma = np.percentile(img[13000:16000,13000:16000], [1,99.8])  # compute percentiles from smaller crop\n",
        "mi, ma = 0, 255                                                   # use min and max dtype values (suitable here)\n",
        "normalizer = MyNormalizer(mi, ma)\n",
        "\n",
        "\n",
        "np.random.seed(16)\n",
        "lbl_cmap = random_label_cmap()\n",
        "\n",
        "#Load the StarDist model\n",
        "\n",
        "model = StarDist2D(None, name=Prediction_model_name, basedir=Prediction_model_path)\n",
        "\n",
        "\n",
        "for image in os.listdir(Data_folder):\n",
        "  print(\"Performing prediction on: \"+image)\n",
        "\n",
        "  X = imread(Data_folder+\"/\"+image)\n",
        "\n",
        "  print(\"Image dimension \"+str(X.shape))\n",
        "\n",
        "  short_name = os.path.splitext(image)\n",
        "\n",
        "  n_channel = 1 if X.ndim == 2 else X.shape[-1]\n",
        "  \n",
        "  # axis_norm = (0,1,2) # normalize channels jointly\n",
        "  if n_channel == 1:\n",
        "    axis_norm = (0,1)   # normalize channels independently\n",
        "    print(\"Normalizing image channels independently\")\n",
        "    block_size = (block_size_Y, block_size_X)\n",
        "    min_overlap = (min_overlap, min_overlap)\n",
        "    n_tiles = (n_tiles_Y, n_tiles_X)\n",
        "    axes=\"YX\"\n",
        "\n",
        "  if n_channel > 1:\n",
        "    axis_norm = (0,1,2) # normalize channels jointly\n",
        "    print(\"Normalizing image channels jointly\")\n",
        "    axes=\"YXC\"\n",
        "    block_size = (block_size_Y, block_size_X, 3)\n",
        "    n_tiles = (n_tiles_Y, n_tiles_X, 1)\n",
        "    min_overlap = (min_overlap, min_overlap, 0)  \n",
        "    sys.stdout.flush()\n",
        "  \n",
        "  zarr.save_array(str(Temp_folder+\"/image.zarr\"), X)\n",
        "  del X\n",
        "  img = zarr.open(str(Temp_folder+\"/image.zarr\"), mode='r')\n",
        "  \n",
        "  labels = zarr.open(str(Temp_folder+\"/labels.zarr\"), mode='w', shape=img.shape[:3], chunks=img.chunks[:3], dtype=np.int32)\n",
        "    \n",
        "  \n",
        "  labels, polygons = model.predict_instances_big(img, axes=axes, block_size=block_size, min_overlap=min_overlap, context=None,\n",
        "                                            normalizer=normalizer, show_progress=True, n_tiles=n_tiles)\n",
        "  \n",
        "# Save the predicted mask in the result folder\n",
        "  os.chdir(Results_folder)\n",
        "  if Mask_images:\n",
        "    imsave(str(short_name[0])+\".tif\", labels)\n",
        "  if Region_of_interests:  \n",
        "    export_imagej_rois(str(short_name[0])+'labels_roi.zip', polygons['coord'], compression=ZIP_DEFLATED)\n",
        "\n",
        "  del img\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",
        "# One example image \n",
        "\n",
        "fig, (a,b) = plt.subplots(1,2, figsize=(20,20))\n",
        "a.imshow(labels[::8,::8], cmap='tab20b')\n",
        "b.imshow(labels[::8,::8], cmap=lbl_cmap)\n",
        "a.axis('off'); b.axis('off');\n",
        "None;"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hvkd66PldsXB"
      },
      "source": [
        "## **6.3. 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": "VxGg1fmkFol2"
      },
      "source": [
        "# **7. Version log**\n",
        "---\n",
        "<font size = 4>**v1.19.3**:  \n",
        "\n",
        "*    freezing bioimageio dependency\n",
        "\n",
        "<font size = 4>**v1.19.2**:  \n",
        "\n",
        "*    Pathing bug fixes\n",
        "\n",
        "<font size = 4>**v1.19.1**:  \n",
        "\n",
        "*    Replaced all absolute pathing with relative pathing\n",
        "\n",
        "<font size = 4>**v1.19**:  \n",
        "* `imread` fixed by importing in the pdf generator\n",
        "* bioimage model zoo export for tf2 fixed\n",
        "* Donwload example dataset for training and testing\n",
        "\n",
        "<font size = 4>**v1.18**:  \n",
        "* minor text fixes\n",
        "* PDF export librairy changed from `fpdf` to `fpdf2`\n",
        "\n",
        "<font size = 4>**v1.17**:  \n",
        "* `imwrite` and `imsave` functions have changed the number of parameters\n",
        "\n",
        "<font size = 4>**v1.16.4**:  \n",
        "* Now uses TF 2.8 to ensure GPU access\n",
        "* Bioimage.io export is not working\n",
        "\n",
        "<font size = 4>**v1.16.3**:  \n",
        "\n",
        "\n",
        "* This forces the installation of TF 2.2 but it does not install any additional CUDA version. It ensures the compatibility with Google Colab.\n",
        "* Allows exporting the model in the BioImage Model Zoo compatible format\n",
        "* Allows loading pertained models from the BioImage Model Zoo\n",
        "* `__future__` dependency is imported at the beginning\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UvSlTaH14s3t"
      },
      "source": [
        "# **Thank you for using StarDist 2D!**"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.6.12"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}