Colab_notebooks/U-Net_3D_ZeroCostDL4Mic.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IkSguVy8Xv83"
      },
      "source": [
        "# **U-Net (3D)**\n",
        " ---\n",
        "\n",
        "<font size = 4> The 3D U-Net was first introduced by [Çiçek et al](https://arxiv.org/abs/1606.06650) for learning dense volumetric segmentations from sparsely annotated ground-truth data building upon the original U-Net architecture by [Ronneberger et al](https://arxiv.org/abs/1505.04597). \n",
        "\n",
        "<font size = 4>**This particular implementation allows supervised learning between any two types of 3D image data. If you are interested in image segmentation of 2D datasets, you should use the 2D U-Net notebook instead.**\n",
        "\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 ([ZeroCostDL4Mic](https://github.com/HenriquesLab/ZeroCostDL4Mic/wiki)) jointly developed by the [Jacquemet](https://cellmig.org/) and [Henriques](https://henriqueslab.github.io/) laboratories and created by Daniel Krentzel. The BioImage Model Zoo export was jointly developed by [Estibaliz Gómez de Mariscal](https://github.com/esgomezm) (deepImageJ team).\n",
        "\n",
        "<font size = 4>This notebook is laregly based on the following paper: \n",
        "\n",
        "<font size = 4>[**3D U-Net: Learning Dense Volumetric Segmentation from Sparse Annotation**](https://arxiv.org/pdf/1606.06650.pdf) by Özgün Çiçek *et al.* published on arXiv in 2016\n",
        "\n",
        "<font size = 4>The following two Python libraries play an important role in the notebook: \n",
        "\n",
        "1. <font size = 4>[**Elasticdeform**](https://github.com/gvtulder/elasticdeform)\n",
        " by Gijs van Tulder was used to augment the 3D training data using elastic grid-based deformations as described in the original 3D U-Net paper. \n",
        "\n",
        "2. <font size = 4>[**Tifffile**](https://github.com/cgohlke/tifffile) by Christoph Gohlke is a great library for reading and writing TIFF files. \n",
        "\n",
        "3. <font size = 4>[**Imgaug**](https://github.com/aleju/imgaug) by Alexander Jung *et al.* is an amazing library for image augmentation in machine learning - it is the most complete and extensive image augmentation package I have found to date. \n",
        "\n",
        "<font size = 4>The [example dataset](https://www.epfl.ch/labs/cvlab/data/data-em/) represents a 5x5x5µm section taken from the CA1 hippocampus region of the brain with annotated mitochondria and was acquired by Graham Knott and Marco Cantoni at EPFL.\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 the original paper and relevant Python libraries when using or developing this notebook.**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jWAz2i7RdxUV"
      },
      "source": [
        "# **How to use this notebook?**\n",
        "\n",
        "---\n",
        "\n",
        "<font size = 4>Video describing how to use ZeroCostDL4Mic 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 cells:  \n",
        "\n",
        "<font size = 4>**Text cells** provide information and can be modified by double-clicking the cell. You are currently reading a text cell. You can create a new one by clicking `+ Text`.\n",
        "\n",
        "<font size = 4>**Code cells** contain code which can be modfied by selecting the cell. To execute the cell, move your cursor to the `[]`-symbol on the left side of the cell (a play button should appear). Click it to execute the cell. Once the cell is fully executed, the animation 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>Three tabs are located on the upper left side of the notebook:\n",
        "\n",
        "1. <font size = 4>*Table of contents* contains the structure of the notebook. Click the headers to move quickly between sections.\n",
        "\n",
        "2. <font size = 4>*Code snippets* provides a wide array of example code specific to Google Colab. You can ignore this when using this notebook.\n",
        "\n",
        "3. <font size = 4>*Files* displays the current working directory. We will mount your Google Drive in Section 1.2. so that you can access your files and save them permanently.\n",
        "\n",
        "<font size = 4>**Important:** All uploaded files are purged once the runtime ends.\n",
        "\n",
        "<font size = 4>**Note:** The directory *sample data* in *Files* contains default files. Do not upload anything there!\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 by clicking *File* -> *Save a copy in Drive*.\n",
        "\n",
        "<font size = 4>To **edit a cell**, double click on the text. This will either display the source code (in code cells) or the [markdown](https://colab.research.google.com/notebooks/markdown_guide.ipynb#scrollTo=70pYkR9LiOV0) (in text cells).\n",
        "You can use `#` in code cells to comment out parts of the code. This allows you to keep the original piece of code while not executing it."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gKDLkLWUd-YX"
      },
      "source": [
        "#**0. Before getting started**\n",
        "---\n",
        "\n",
        "<font size = 4>As the network operates in three dimensions, certain consideration should be given to correctly pre-processing the data. Ensure that the structure of interest does not substantially change between slices - image volumes with isotropic pixelsizes are ideal for this architecture.\n",
        "\n",
        "<font size = 4>Each image volume must be provided as an **8-bit** or **binary multipage TIFF file** to maintain the correct ordering of individual image slices. If more than one image volume has been annotated, source and target files must be named identically and placed in separate directories. In case only one image volume has been annotated, source and target file do not have to be placed in separate directories and can be named differently, as long as their paths are explicitly provided in Section 3. \n",
        "\n",
        "<font size = 4>**Prepare two datasets** (*training* and *testing*) for quality control puproses. Make sure that the *testing* dataset does not overlap with the *training* dataset and is ideally sourced from a different acquisiton and sample to ensure robustness of the trained model. \n",
        "\n",
        "\n",
        "---\n",
        "\n",
        "\n",
        "### **Directory structure**\n",
        "\n",
        "<font size = 4>Make sure to adhere to one of the following directory structures. If only one annotated training volume exists, choose the first structure. In case more than one training volume is available, choose the second structure.\n",
        "\n",
        "<font size = 4>**Structure 1:** Only one training volume\n",
        "```\n",
        "path/to/directory/with/one/training/volume\n",
        "│--training_source.tif\n",
        "│--training_target.tif\n",
        "|   \n",
        "│--testing_source.tif\n",
        "|--testing_target.tif \n",
        "|\n",
        "|--data_to_predict_on.tif\n",
        "|--prediction_results.tif\n",
        "\n",
        "```\n",
        "<font size = 4>**Structure 2:** Various training volumes\n",
        "```\n",
        "path/to/directory/with/various/training/volumes\n",
        "│--testing_source.tif\n",
        "|--testing_target.tif \n",
        "|\n",
        "└───training\n",
        "|   └───source\n",
        "|   |   |--training_volume_one.tif\n",
        "|   |   |--training_volume_two.tif\n",
        "|   |   |--...\n",
        "|   |   |--training_volume_n.tif\n",
        "|   |\n",
        "|   └───target\n",
        "|       |--training_volume_one.tif\n",
        "|       |--training_volume_two.tif\n",
        "|       |--...\n",
        "|       |--training_volume_n.tif\n",
        "|\n",
        "|--data_to_predict_on.tif\n",
        "|--prediction_results.tif\n",
        "```\n",
        "<font size = 4>**Note:** Naming directories is completely up to you, as long as the paths are correctly specified throughout the notebook.\n",
        "\n",
        "\n",
        "---\n",
        "\n",
        "\n",
        "### **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 so), 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 3D U-Net dependencies**\n",
        "---\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PVbFEzo1DgUB"
      },
      "source": [
        "## **1.1. Install key dependencies**\n",
        "---\n",
        "<font size = 4> "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "q4wM9Sr0Dbbf"
      },
      "outputs": [],
      "source": [
        "#@markdown ##Play to install U-Net dependencies\n",
        "# Install packages which are not included in Google Colab\n",
        "\n",
        "!pip install data\n",
        "!pip install -q tifffile\n",
        "!pip install -q wget\n",
        "!pip install -q fpdf2\n",
        "!pip install -q PTable # Nice tables \n",
        "!pip install -q zarr\n",
        "!pip install -q imagecodecs\n",
        "\n",
        "!pip install \"bioimageio.core>=0.5\"\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3pcnkKZaDjaF"
      },
      "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>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9rNP7LpoDoGk"
      },
      "source": [
        "## **1.3. Load key dependencies**\n",
        "---\n",
        "<font size = 4> "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "fq21zJVFNASx"
      },
      "outputs": [],
      "source": [
        "from __future__ import absolute_import, division, print_function, unicode_literals\n",
        "\n",
        "#@markdown ##Load key 3D U-Net dependencies and instantiate network\n",
        "Notebook_version = '2.1'\n",
        "Network = 'U-Net (3D)'\n",
        "\n",
        "from builtins import any as b_any\n",
        "\n",
        "def get_requirements_path():\n",
        "    # Store requirements file in 'contents' directory \n",
        "    current_dir = os.getcwd()\n",
        "    dir_count = current_dir.count('/') - 1\n",
        "    path = '../' * (dir_count) + 'requirements.txt'\n",
        "    return path\n",
        "\n",
        "def filter_files(file_list, filter_list):\n",
        "    filtered_list = []\n",
        "    for fname in file_list:\n",
        "        if b_any(fname.split('==')[0] in s for s in filter_list):\n",
        "            filtered_list.append(fname)\n",
        "    return filtered_list\n",
        "\n",
        "def build_requirements_file(before, after):\n",
        "    path = get_requirements_path()\n",
        "\n",
        "    # Exporting requirements.txt for local run\n",
        "    !pip freeze > $path\n",
        "\n",
        "    # Get minimum requirements file\n",
        "    df = pd.read_csv(path)\n",
        "    mod_list = [m.split('.')[0] for m in after if not m in before]\n",
        "    req_list_temp = df.values.tolist()\n",
        "    req_list = [x[0] for x in req_list_temp]\n",
        "\n",
        "    # Replace with package name and handle cases where import name is different to module name\n",
        "    mod_name_list = [['sklearn', 'scikit-learn'], ['skimage', 'scikit-image']]\n",
        "    mod_replace_list = [[x[1] for x in mod_name_list] if s in [x[0] for x in mod_name_list] else s for s in mod_list] \n",
        "    filtered_list = filter_files(req_list, mod_replace_list)\n",
        "\n",
        "    file=open(path,'w')\n",
        "    for item in filtered_list:\n",
        "        file.writelines(item)\n",
        "\n",
        "    file.close()\n",
        "\n",
        "import sys\n",
        "before = [str(m) for m in sys.modules]\n",
        "\n",
        "#Put the imported code and libraries here\n",
        "# !pip install fpdf\n",
        "\n",
        "try:\n",
        "    import elasticdeform\n",
        "except:\n",
        "    !pip install elasticdeform\n",
        "    import elasticdeform\n",
        "\n",
        "try:\n",
        "    import tifffile\n",
        "except:\n",
        "    !pip install tifffile\n",
        "    import tifffile\n",
        "\n",
        "try:\n",
        "    import imgaug.augmenters as iaa\n",
        "except:\n",
        "    !pip install imgaug\n",
        "    import imgaug.augmenters as iaa\n",
        "\n",
        "import os\n",
        "import csv\n",
        "import random\n",
        "import h5py\n",
        "import imageio\n",
        "import math\n",
        "import shutil\n",
        "\n",
        "import pandas as pd\n",
        "from glob import glob\n",
        "from tqdm import tqdm\n",
        "\n",
        "from skimage import transform\n",
        "from skimage import exposure\n",
        "from skimage import color\n",
        "from skimage import io\n",
        "\n",
        "from scipy.ndimage import zoom\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "import numpy as np\n",
        "\n",
        "import tensorflow as tf\n",
        "print(\"TensorFlow version: {}\".format(tf.__version__))\n",
        "\n",
        "# from keras import backend as K\n",
        "\n",
        "# from keras.layers import Conv3D\n",
        "# from keras.layers import BatchNormalization\n",
        "# from keras.layers import ReLU\n",
        "# from keras.layers import MaxPooling3D\n",
        "# from keras.layers import Conv3DTranspose\n",
        "# from keras.layers import Input\n",
        "# from keras.layers import Concatenate\n",
        "\n",
        "# from keras.models import Model\n",
        "\n",
        "# from keras.utils import Sequence\n",
        "# from keras.callbacks import ModelCheckpoint\n",
        "# from keras.callbacks import CSVLogger\n",
        "# from keras.callbacks import Callback\n",
        "\n",
        "from tensorflow.keras import backend as K\n",
        "\n",
        "from tensorflow.keras.layers import Conv3D\n",
        "from tensorflow.keras.layers import BatchNormalization\n",
        "from tensorflow.keras.layers import ReLU\n",
        "from tensorflow.keras.layers import MaxPooling3D\n",
        "from tensorflow.keras.layers import Conv3DTranspose\n",
        "from tensorflow.keras.layers import Input\n",
        "from tensorflow.keras.layers import Concatenate\n",
        "\n",
        "from tensorflow.keras.models import Model\n",
        "\n",
        "from tensorflow.keras.utils import Sequence\n",
        "from tensorflow.keras.callbacks import ModelCheckpoint\n",
        "from tensorflow.keras.callbacks import CSVLogger\n",
        "from tensorflow.keras.callbacks import Callback\n",
        "\n",
        "from tensorflow.keras.metrics import RootMeanSquaredError\n",
        "\n",
        "from tensorflow.keras.optimizers import Adam, SGD, RMSprop\n",
        "\n",
        "from ipywidgets import interact\n",
        "from ipywidgets import interactive\n",
        "from ipywidgets import fixed\n",
        "from ipywidgets import interact_manual \n",
        "import ipywidgets as widgets\n",
        "\n",
        "from fpdf import FPDF, HTMLMixin\n",
        "from datetime import datetime\n",
        "import subprocess\n",
        "from pip._internal.operations.freeze import freeze\n",
        "import time\n",
        "\n",
        "from skimage import io\n",
        "import matplotlib\n",
        "\n",
        "from skimage import io\n",
        "from shutil import rmtree\n",
        "from bioimageio.core.build_spec import build_model, add_weights\n",
        "from bioimageio.core.resource_tests import test_model\n",
        "from bioimageio.core.weight_converter.keras import convert_weights_to_tensorflow_saved_model_bundle\n",
        "\n",
        "print(\"Dependencies installed and imported.\")\n",
        "\n",
        "# Define MultiPageTiffGenerator class\n",
        "class MultiPageTiffGenerator(Sequence):\n",
        "\n",
        "    def __init__(self,\n",
        "                 source_path,\n",
        "                 target_path,\n",
        "                 batch_size=1,\n",
        "                 shape=(128,128,32,1),\n",
        "                 augment=False,\n",
        "                 augmentations=[],\n",
        "                 deform_augment=False,\n",
        "                 deform_augmentation_params=(5,3,4),\n",
        "                 val_split=0.2,\n",
        "                 is_val=False,\n",
        "                 random_crop=True,\n",
        "                 downscale=1,\n",
        "                 binary_target=False):\n",
        "\n",
        "        # If directory with various multi-page tiffiles is provided read as list\n",
        "        if os.path.isfile(source_path):\n",
        "            self.dir_flag = False\n",
        "            self.source = tifffile.imread(source_path)\n",
        "            if binary_target:\n",
        "                self.target = tifffile.imread(target_path).astype(bool)\n",
        "            else:\n",
        "                self.target = tifffile.imread(target_path)\n",
        "\n",
        "        elif os.path.isdir(source_path):\n",
        "            self.dir_flag = True\n",
        "            self.source_dir_list = glob(os.path.join(source_path, '*'))\n",
        "            self.target_dir_list = glob(os.path.join(target_path, '*'))\n",
        "\n",
        "            self.source_dir_list.sort()\n",
        "            self.target_dir_list.sort()\n",
        "\n",
        "        self.shape = shape\n",
        "        self.batch_size = batch_size\n",
        "        self.augment = augment\n",
        "        self.val_split = val_split\n",
        "        self.is_val = is_val\n",
        "        self.random_crop = random_crop\n",
        "        self.downscale = downscale\n",
        "        self.binary_target = binary_target\n",
        "        self.deform_augment = deform_augment\n",
        "        self.on_epoch_end()\n",
        "        \n",
        "        if self.augment:\n",
        "            # pass list of augmentation functions \n",
        "            self.seq = iaa.Sequential(augmentations, random_order=True) # apply augmenters in random order\n",
        "        if self.deform_augment:\n",
        "            self.deform_sigma, self.deform_points, self.deform_order = deform_augmentation_params\n",
        "\n",
        "    def __len__(self):\n",
        "        # If various multi-page tiff files provided sum all images within each\n",
        "        if self.augment:\n",
        "            augment_factor = 4\n",
        "        else:\n",
        "            augment_factor = 1\n",
        "    \n",
        "        if self.dir_flag:\n",
        "            num_of_imgs = 0\n",
        "            for tiff_path in self.source_dir_list:\n",
        "                num_of_imgs += tifffile.imread(tiff_path).shape[0]\n",
        "            xy_shape = tifffile.imread(self.source_dir_list[0]).shape[1:]\n",
        "\n",
        "            if self.is_val:\n",
        "                if self.random_crop:\n",
        "                    crop_volume = self.shape[0] * self.shape[1] * self.shape[2]\n",
        "                    volume = xy_shape[0] * xy_shape[1] * self.val_split * num_of_imgs\n",
        "                    return math.floor(augment_factor * volume / (crop_volume * self.batch_size * self.downscale))\n",
        "                else:\n",
        "                    return math.floor(self.val_split * num_of_imgs / self.batch_size)\n",
        "            else:\n",
        "                if self.random_crop:\n",
        "                    crop_volume = self.shape[0] * self.shape[1] * self.shape[2]\n",
        "                    volume = xy_shape[0] * xy_shape[1] * (1 - self.val_split) * num_of_imgs\n",
        "                    return math.floor(augment_factor * volume / (crop_volume * self.batch_size * self.downscale))\n",
        "\n",
        "                else:\n",
        "                    return math.floor(augment_factor*(1 - self.val_split) * num_of_imgs/self.batch_size)\n",
        "        else:\n",
        "            if self.is_val:\n",
        "                if self.random_crop:\n",
        "                    crop_volume = self.shape[0] * self.shape[1] * self.shape[2]\n",
        "                    volume = self.source.shape[0] * self.source.shape[1] * self.val_split * self.source.shape[2]\n",
        "                    return math.floor(augment_factor * volume / (crop_volume * self.batch_size * self.downscale))\n",
        "                else:\n",
        "                    return math.floor((self.val_split * self.source.shape[0] / self.batch_size))\n",
        "            else:\n",
        "                if self.random_crop:\n",
        "                    crop_volume = self.shape[0] * self.shape[1] * self.shape[2]\n",
        "                    volume = self.source.shape[0] * self.source.shape[1] * (1 - self.val_split) * self.source.shape[2]\n",
        "                    return math.floor(augment_factor * volume / (crop_volume * self.batch_size * self.downscale))\n",
        "                else:\n",
        "                    return math.floor(augment_factor * (1 - self.val_split) * self.source.shape[0] / self.batch_size)\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        source_batch = np.empty((self.batch_size,\n",
        "                                 self.shape[0],\n",
        "                                 self.shape[1],\n",
        "                                 self.shape[2],\n",
        "                                 self.shape[3]))\n",
        "        target_batch = np.empty((self.batch_size,\n",
        "                                 self.shape[0],\n",
        "                                 self.shape[1],\n",
        "                                 self.shape[2],\n",
        "                                 self.shape[3]))\n",
        "\n",
        "        for batch in range(self.batch_size):\n",
        "            # Modulo operator ensures IndexError is avoided\n",
        "            stack_start = self.batch_list[(idx+batch*self.shape[2])%len(self.batch_list)]\n",
        "\n",
        "            if self.dir_flag:\n",
        "                self.source = tifffile.imread(self.source_dir_list[stack_start[0]])\n",
        "                if self.binary_target:\n",
        "                    self.target = tifffile.imread(self.target_dir_list[stack_start[0]]).astype(bool)\n",
        "                else:\n",
        "                    self.target = tifffile.imread(self.target_dir_list[stack_start[0]])\n",
        "\n",
        "            src_list = []\n",
        "            tgt_list = []\n",
        "            for i in range(stack_start[1], stack_start[1]+self.shape[2]):\n",
        "                src = self.source[i]\n",
        "                src = transform.downscale_local_mean(src, (self.downscale, self.downscale))\n",
        "                if not self.random_crop:\n",
        "                    src = transform.resize(src, (self.shape[0], self.shape[1]), mode='constant', preserve_range=True)\n",
        "                src = self._min_max_scaling(src)\n",
        "                src_list.append(src)\n",
        "\n",
        "                tgt = self.target[i]\n",
        "                tgt = transform.downscale_local_mean(tgt, (self.downscale, self.downscale))\n",
        "                if not self.random_crop:\n",
        "                    tgt = transform.resize(tgt, (self.shape[0], self.shape[1]), mode='constant', preserve_range=True)\n",
        "                if not self.binary_target:\n",
        "                    tgt = self._min_max_scaling(tgt)\n",
        "                tgt_list.append(tgt)\n",
        "\n",
        "            if self.random_crop:\n",
        "                if src.shape[0] == self.shape[0]:\n",
        "                    x_rand = 0\n",
        "                if src.shape[1] == self.shape[1]:\n",
        "                    y_rand = 0\n",
        "                if src.shape[0] > self.shape[0]:\n",
        "                    x_rand = np.random.randint(src.shape[0] - self.shape[0])\n",
        "                if src.shape[1] > self.shape[1]:\n",
        "                    y_rand = np.random.randint(src.shape[1] - self.shape[1])\n",
        "                if src.shape[0] < self.shape[0] or src.shape[1] < self.shape[1]:\n",
        "                    raise ValueError('Patch shape larger than (downscaled) source shape')\n",
        "            \n",
        "            for i in range(self.shape[2]):\n",
        "                if self.random_crop:\n",
        "                    src = src_list[i]\n",
        "                    tgt = tgt_list[i]\n",
        "                    src_crop = src[x_rand:self.shape[0]+x_rand, y_rand:self.shape[1]+y_rand]\n",
        "                    tgt_crop = tgt[x_rand:self.shape[0]+x_rand, y_rand:self.shape[1]+y_rand]\n",
        "                else:\n",
        "                    src_crop = src_list[i]\n",
        "                    tgt_crop = tgt_list[i]\n",
        "\n",
        "                source_batch[batch,:,:,i,0] = src_crop\n",
        "                target_batch[batch,:,:,i,0] = tgt_crop\n",
        "\n",
        "        if self.augment:\n",
        "            # On-the-fly data augmentation\n",
        "            source_batch, target_batch = self.augment_volume(source_batch, target_batch)\n",
        "\n",
        "            # Data augmentation by reversing stack\n",
        "            if np.random.random() > 0.5:\n",
        "                source_batch, target_batch = source_batch[::-1], target_batch[::-1]\n",
        "            \n",
        "            # Data augmentation by elastic deformation\n",
        "            if np.random.random() > 0.5 and self.deform_augment:\n",
        "                source_batch, target_batch = self.deform_volume(source_batch, target_batch)\n",
        "            \n",
        "            if not self.binary_target:\n",
        "                target_batch = self._min_max_scaling(target_batch)\n",
        "            \n",
        "            return self._min_max_scaling(source_batch), target_batch\n",
        "        \n",
        "        else:\n",
        "            return source_batch, target_batch\n",
        "\n",
        "    def on_epoch_end(self):\n",
        "        # Validation split performed here\n",
        "        self.batch_list = []\n",
        "        # Create batch_list of all combinations of tifffile and stack position\n",
        "        if self.dir_flag:\n",
        "            for i in range(len(self.source_dir_list)):\n",
        "                num_of_pages = tifffile.imread(self.source_dir_list[i]).shape[0]\n",
        "                if self.is_val:\n",
        "                    start_page = num_of_pages-math.floor(self.val_split*num_of_pages)\n",
        "                    for j in range(start_page, num_of_pages-self.shape[2]):\n",
        "                      self.batch_list.append([i, j])\n",
        "                else:\n",
        "                    last_page = math.floor((1-self.val_split)*num_of_pages)\n",
        "                    for j in range(last_page-self.shape[2]):\n",
        "                        self.batch_list.append([i, j])\n",
        "        else:\n",
        "            num_of_pages = self.source.shape[0]\n",
        "            if self.is_val:\n",
        "                start_page = num_of_pages-math.floor(self.val_split*num_of_pages)\n",
        "                for j in range(start_page, num_of_pages-self.shape[2]):\n",
        "                    self.batch_list.append([0, j])\n",
        "\n",
        "            else:\n",
        "                last_page = math.floor((1-self.val_split)*num_of_pages)\n",
        "                for j in range(last_page-self.shape[2]):\n",
        "                    self.batch_list.append([0, j])\n",
        "        \n",
        "        if self.is_val and (len(self.batch_list) <= 0):\n",
        "            raise ValueError('validation_split too small! Increase val_split or decrease z-depth')\n",
        "        random.shuffle(self.batch_list)\n",
        "   \n",
        "    def _min_max_scaling(self, data):\n",
        "        n = data - np.min(data)\n",
        "        d = np.max(data) - np.min(data) \n",
        "        \n",
        "        return n/d\n",
        "   \n",
        "    def class_weights(self):\n",
        "        ones = 0\n",
        "        pixels = 0\n",
        "\n",
        "        if self.dir_flag:\n",
        "            for i in range(len(self.target_dir_list)):\n",
        "                tgt = tifffile.imread(self.target_dir_list[i]).astype(bool)\n",
        "                ones += np.sum(tgt)\n",
        "                pixels += tgt.shape[0]*tgt.shape[1]*tgt.shape[2]\n",
        "        else:\n",
        "          ones = np.sum(self.target)\n",
        "          pixels = self.target.shape[0]*self.target.shape[1]*self.target.shape[2]\n",
        "        p_ones = ones/pixels\n",
        "        p_zeros = 1-p_ones\n",
        "\n",
        "        # Return swapped probability to increase weight of unlikely class\n",
        "        return p_ones, p_zeros\n",
        "\n",
        "    def deform_volume(self, src_vol, tgt_vol):\n",
        "        [src_dfrm, tgt_dfrm] = elasticdeform.deform_random_grid([src_vol, tgt_vol],\n",
        "                                                                axis=(1, 2, 3),\n",
        "                                                                sigma=self.deform_sigma,\n",
        "                                                                points=self.deform_points,\n",
        "                                                                order=self.deform_order)\n",
        "        if self.binary_target:\n",
        "            tgt_dfrm = tgt_dfrm > 0.1\n",
        "        \n",
        "        return self._min_max_scaling(src_dfrm), tgt_dfrm \n",
        "\n",
        "    def augment_volume(self, src_vol, tgt_vol):\n",
        "        \n",
        "        src_vol_aug = np.empty(src_vol.shape)\n",
        "        tgt_vol_aug = np.empty(tgt_vol.shape)\n",
        "\n",
        "        for i in range(src_vol.shape[3]):\n",
        "            src_aug_z, tgt_aug_z = self.seq(images=src_vol[:,:,:,i,0].astype('float16'), \n",
        "                                                                      segmentation_maps=np.expand_dims(tgt_vol[:,:,:,i,0].astype(bool), axis=-1))\n",
        "            src_vol_aug[:,:,:,i,0] = src_aug_z \n",
        "            tgt_vol_aug[:,:,:,i,0] = np.squeeze(tgt_aug_z)\n",
        "        return self._min_max_scaling(src_vol_aug), tgt_vol_aug\n",
        "\n",
        "    def sample_augmentation(self, idx):\n",
        "        src, tgt = self.__getitem__(idx)\n",
        "\n",
        "        src_aug, tgt_aug = self.augment_volume(src, tgt)\n",
        "        \n",
        "        if self.deform_augment:\n",
        "            src_aug, tgt_aug = self.deform_volume(src_aug, tgt_aug)\n",
        "\n",
        "        return src_aug, tgt_aug \n",
        "\n",
        "# Define custom loss and dice coefficient\n",
        "def dice_coefficient(y_true, y_pred):\n",
        "    eps = 1e-6\n",
        "    y_true_f = K.flatten(y_true)\n",
        "    y_pred_f = K.flatten(y_pred)\n",
        "    intersection = K.sum(y_true_f*y_pred_f)\n",
        "\n",
        "    return (2.*intersection)/(K.sum(y_true_f*y_true_f)+K.sum(y_pred_f*y_pred_f)+eps)\n",
        "\n",
        "def weighted_binary_crossentropy(zero_weight, one_weight):\n",
        "    def _weighted_binary_crossentropy(y_true, y_pred):\n",
        "        binary_crossentropy = K.binary_crossentropy(y_true, y_pred)\n",
        "\n",
        "        weight_vector = y_true*one_weight+(1.-y_true)*zero_weight\n",
        "        weighted_binary_crossentropy = weight_vector*binary_crossentropy\n",
        "\n",
        "        return K.mean(weighted_binary_crossentropy)\n",
        "\n",
        "    return _weighted_binary_crossentropy\n",
        "\n",
        "# Custom callback showing sample prediction\n",
        "class SampleImageCallback(Callback):\n",
        "\n",
        "    def __init__(self, model, sample_data, model_path, save=False):\n",
        "        self.model = model\n",
        "        self.sample_data = sample_data\n",
        "        self.model_path = model_path\n",
        "        self.save = save\n",
        "\n",
        "    def on_epoch_end(self, epoch, logs={}):\n",
        "        sample_predict = self.model.predict_on_batch(self.sample_data)\n",
        "\n",
        "        f=plt.figure(figsize=(16,8))\n",
        "        plt.subplot(1,2,1)\n",
        "        plt.imshow(self.sample_data[0,:,:,0,0], interpolation='nearest', cmap='gray')\n",
        "        plt.title('Sample source')\n",
        "        plt.axis('off');\n",
        "\n",
        "        plt.subplot(1,2,2)\n",
        "        plt.imshow(sample_predict[0,:,:,0,0], interpolation='nearest', cmap='magma')\n",
        "        plt.title('Predicted target')\n",
        "        plt.axis('off');\n",
        "\n",
        "        plt.show()\n",
        "\n",
        "        if self.save:\n",
        "            plt.savefig(self.model_path + '/epoch_' + str(epoch+1) + '.png')\n",
        "\n",
        "\n",
        "# Define Unet3D class\n",
        "class Unet3D:\n",
        "\n",
        "    def __init__(self,\n",
        "                 shape=(256,256,16,1)):\n",
        "        if isinstance(shape, str):\n",
        "            shape = eval(shape)\n",
        "\n",
        "        self.shape = shape\n",
        "        \n",
        "        input_tensor = Input(self.shape, name='input')\n",
        "\n",
        "        self.model = self.unet_3D(input_tensor)\n",
        "\n",
        "    def down_block_3D(self, input_tensor, filters):\n",
        "        x = Conv3D(filters=filters, kernel_size=(3,3,3), padding='same')(input_tensor)\n",
        "        x = BatchNormalization()(x)\n",
        "        x = ReLU()(x)\n",
        "\n",
        "        x = Conv3D(filters=filters*2, kernel_size=(3,3,3), padding='same')(x)\n",
        "        x = BatchNormalization()(x)\n",
        "        x = ReLU()(x)\n",
        "\n",
        "        return x\n",
        "\n",
        "    def up_block_3D(self, input_tensor, concat_layer, filters):\n",
        "        x = Conv3DTranspose(filters, kernel_size=(2,2,2), strides=(2,2,2))(input_tensor)\n",
        "\n",
        "        x = Concatenate()([x, concat_layer])\n",
        "\n",
        "        x = Conv3D(filters=filters, kernel_size=(3,3,3), padding='same')(x)\n",
        "        x = BatchNormalization()(x)\n",
        "        x = ReLU()(x)\n",
        "\n",
        "        x = Conv3D(filters=filters*2, kernel_size=(3,3,3), padding='same')(x)\n",
        "        x = BatchNormalization()(x)\n",
        "        x = ReLU()(x)\n",
        "\n",
        "        return x\n",
        "\n",
        "    def unet_3D(self, input_tensor, filters=32):\n",
        "        d1 = self.down_block_3D(input_tensor, filters=filters)\n",
        "        p1 = MaxPooling3D(pool_size=(2,2,2), strides=(2,2,2), data_format='channels_last')(d1)\n",
        "        d2 = self.down_block_3D(p1, filters=filters*2)\n",
        "        p2 = MaxPooling3D(pool_size=(2,2,2), strides=(2,2,2), data_format='channels_last')(d2)\n",
        "        d3 = self.down_block_3D(p2, filters=filters*4)\n",
        "        p3 = MaxPooling3D(pool_size=(2,2,2), strides=(2,2,2), data_format='channels_last')(d3)\n",
        "\n",
        "        d4 = self.down_block_3D(p3, filters=filters*8)\n",
        "\n",
        "        u1 = self.up_block_3D(d4, d3, filters=filters*4)\n",
        "        u2 = self.up_block_3D(u1, d2, filters=filters*2)\n",
        "        u3 = self.up_block_3D(u2, d1, filters=filters)\n",
        "\n",
        "        output_tensor = Conv3D(filters=1, kernel_size=(1,1,1), activation='sigmoid')(u3)\n",
        "\n",
        "        return Model(inputs=[input_tensor], outputs=[output_tensor])\n",
        "\n",
        "    def summary(self):\n",
        "        return self.model.summary()\n",
        "\n",
        "    # Pass generators instead\n",
        "    def train(self, \n",
        "              epochs, \n",
        "              batch_size, \n",
        "              train_generator,\n",
        "              val_generator, \n",
        "              model_path, \n",
        "              model_name,\n",
        "              optimizer='adam',\n",
        "              learning_rate=0.001,\n",
        "              loss='weighted_binary_crossentropy',\n",
        "              metrics='dice',\n",
        "              ckpt_period=1, \n",
        "              save_best_ckpt_only=False, \n",
        "              ckpt_path=None):\n",
        "\n",
        "        class_weight_zero, class_weight_one = train_generator.class_weights()\n",
        "        \n",
        "        if loss == 'weighted_binary_crossentropy':\n",
        "            loss = weighted_binary_crossentropy(class_weight_zero, class_weight_one)\n",
        "        \n",
        "        if metrics == 'dice':\n",
        "            metrics = dice_coefficient\n",
        "\n",
        "        if optimizer == 'adam':\n",
        "            optimizer = Adam(learning_rate=learning_rate)\n",
        "        elif optimizer == 'sgd':\n",
        "            optimizer = SGD(learning_rate=learning_rate)\n",
        "        elif optimizer == 'rmsprop':\n",
        "            optimizer = RMSprop(learning_rate=learning_rate)\n",
        "\n",
        "        self.model.compile(optimizer=optimizer,\n",
        "                           loss=loss,\n",
        "                           metrics=[metrics])\n",
        "\n",
        "        if ckpt_path is not None:\n",
        "            self.model.load_weights(ckpt_path)\n",
        "\n",
        "        full_model_path = os.path.join(model_path, model_name)\n",
        "\n",
        "        if not os.path.exists(full_model_path):\n",
        "            os.makedirs(full_model_path)\n",
        "        \n",
        "        log_dir = full_model_path + '/Quality Control'\n",
        "\n",
        "        if not os.path.exists(log_dir):\n",
        "            os.makedirs(log_dir)\n",
        "        \n",
        "        ckpt_dir =  full_model_path + '/ckpt'\n",
        "\n",
        "        if not os.path.exists(ckpt_dir):\n",
        "            os.makedirs(ckpt_dir)\n",
        "\n",
        "        csv_out_name = log_dir + '/training_evaluation.csv'\n",
        "        if ckpt_path is None:\n",
        "            csv_logger = CSVLogger(csv_out_name)\n",
        "        else:\n",
        "            csv_logger = CSVLogger(csv_out_name, append=True)\n",
        "\n",
        "        if save_best_ckpt_only:\n",
        "            ckpt_name = ckpt_dir + '/' + model_name + '.hdf5'\n",
        "        else:\n",
        "            ckpt_name = ckpt_dir + '/' + model_name + '_epoch_{epoch:02d}_val_loss_{val_loss:.4f}.hdf5'\n",
        "        \n",
        "        model_ckpt = ModelCheckpoint(ckpt_name,\n",
        "                                     verbose=1,\n",
        "                                     save_freq=ckpt_period,\n",
        "                                     save_best_only=save_best_ckpt_only,\n",
        "                                     save_weights_only=True)\n",
        "\n",
        "        sample_batch, __ = val_generator.__getitem__(random.randint(0, len(val_generator)))\n",
        "        sample_img = SampleImageCallback(self.model, \n",
        "                                         sample_batch, \n",
        "                                         model_path)\n",
        "\n",
        "        self.model.fit(train_generator,\n",
        "                       validation_data=val_generator,\n",
        "                       validation_steps=math.floor(len(val_generator)/batch_size),\n",
        "                       epochs=epochs,\n",
        "                       callbacks=[csv_logger,\n",
        "                                  model_ckpt,\n",
        "                                  sample_img])\n",
        "\n",
        "        last_ckpt_name = ckpt_dir + '/' + model_name + '_last.hdf5'\n",
        "        self.model.save_weights(last_ckpt_name)\n",
        "\n",
        "    def _min_max_scaling(self, data):\n",
        "        n = data - np.min(data)\n",
        "        d = np.max(data) - np.min(data) \n",
        "        \n",
        "        return n/d\n",
        "\n",
        "    def predict(self, \n",
        "                input, \n",
        "                ckpt_path, \n",
        "                z_range=None, \n",
        "                downscaling=None, \n",
        "                true_patch_size=None):\n",
        "\n",
        "        self.model.load_weights(ckpt_path)\n",
        "\n",
        "        if isinstance(downscaling, str):\n",
        "            downscaling = eval(downscaling)\n",
        "\n",
        "        if math.isnan(downscaling):\n",
        "            downscaling = None\n",
        "\n",
        "        if isinstance(true_patch_size, str):\n",
        "            true_patch_size = eval(true_patch_size)\n",
        "        \n",
        "        if not isinstance(true_patch_size, tuple): \n",
        "            if math.isnan(true_patch_size):\n",
        "                true_patch_size = None\n",
        "\n",
        "        if isinstance(input, str):\n",
        "            src_volume = tifffile.imread(input)\n",
        "        elif isinstance(input, np.ndarray):\n",
        "            src_volume = input\n",
        "        else:\n",
        "            raise TypeError('Input is not path or numpy array!')\n",
        "        \n",
        "        in_size = src_volume.shape\n",
        "\n",
        "        if downscaling or true_patch_size is not None:\n",
        "            x_scaling = 0\n",
        "            y_scaling = 0\n",
        "\n",
        "            if true_patch_size is not None:\n",
        "                x_scaling += true_patch_size[0]/self.shape[0]\n",
        "                y_scaling += true_patch_size[1]/self.shape[1]\n",
        "            if downscaling is not None:\n",
        "                x_scaling += downscaling\n",
        "                y_scaling += downscaling\n",
        "\n",
        "            src_list = []\n",
        "            for i in range(src_volume.shape[0]):\n",
        "                 src_list.append(transform.downscale_local_mean(src_volume[i], (int(x_scaling), int(y_scaling))))\n",
        "            src_volume = np.array(src_list)          \n",
        "\n",
        "        if z_range is not None:\n",
        "            src_volume = src_volume[z_range[0]:z_range[1]]\n",
        "\n",
        "        src_volume = self._min_max_scaling(src_volume)       \n",
        "\n",
        "        src_array = np.zeros((1,\n",
        "                              math.ceil(src_volume.shape[1]/self.shape[0])*self.shape[0], \n",
        "                              math.ceil(src_volume.shape[2]/self.shape[1])*self.shape[1],\n",
        "                              math.ceil(src_volume.shape[0]/self.shape[2])*self.shape[2], \n",
        "                              self.shape[3]))\n",
        "\n",
        "        for i in range(src_volume.shape[0]):\n",
        "            src_array[0,:src_volume.shape[1],:src_volume.shape[2],i,0] = src_volume[i]\n",
        "\n",
        "        pred_array = np.empty(src_array.shape)\n",
        "        print(src_volume.dtype)\n",
        "        for i in range(math.ceil(src_volume.shape[1]/self.shape[0])):\n",
        "          for j in range(math.ceil(src_volume.shape[2]/self.shape[1])):\n",
        "            for k in range(math.ceil(src_volume.shape[0]/self.shape[2])):\n",
        "                pred_temp = self.model.predict(src_array[:,\n",
        "                                                         i*self.shape[0]:i*self.shape[0]+self.shape[0],\n",
        "                                                         j*self.shape[1]:j*self.shape[1]+self.shape[1],\n",
        "                                                         k*self.shape[2]:k*self.shape[2]+self.shape[2]])\n",
        "                pred_array[:,\n",
        "                           i*self.shape[0]:i*self.shape[0]+self.shape[0],\n",
        "                           j*self.shape[1]:j*self.shape[1]+self.shape[1],\n",
        "                           k*self.shape[2]:k*self.shape[2]+self.shape[2]] = pred_temp\n",
        "                           \n",
        "        pred_volume = np.rollaxis(np.squeeze(pred_array), -1)[:src_volume.shape[0],:src_volume.shape[1],:src_volume.shape[2]]            \n",
        "\n",
        "        if downscaling is not None:\n",
        "            pred_list = []\n",
        "            for i in range(pred_volume.shape[0]):\n",
        "                 pred_list.append(transform.resize(pred_volume[i], (in_size[1], in_size[2]), preserve_range=True))\n",
        "            pred_volume = np.array(pred_list)\n",
        "\n",
        "        return pred_volume\n",
        "\n",
        "def pdf_export(trained = False, augmentation = False, pretrained_model = False):\n",
        "  class MyFPDF(FPDF, HTMLMixin):\n",
        "    pass\n",
        "\n",
        "  pdf = MyFPDF()\n",
        "  pdf.add_page()\n",
        "  pdf.set_right_margin(-1)\n",
        "  pdf.set_font(\"Arial\", size = 11, style='B') \n",
        "\n",
        "  day = datetime.now()\n",
        "  datetime_str = str(day)[0:10]\n",
        "\n",
        "  Header = 'Training report for '+Network+' model ('+model_name+')\\nDate: '+datetime_str\n",
        "  pdf.multi_cell(180, 5, txt = Header, align = 'L') \n",
        "  pdf.ln(1)\n",
        "\n",
        "  # add another cell \n",
        "  if trained:\n",
        "    training_time = \"Training time: \"+str(hour)+ \"hour(s) \"+str(mins)+\"min(s) \"+str(round(sec))+\"sec(s)\"\n",
        "    pdf.cell(190, 5, txt = training_time, ln = 1, align='L')\n",
        "  pdf.ln(1)\n",
        "\n",
        "  Header_2 = 'Information for your materials and methods:'\n",
        "  pdf.cell(190, 5, txt=Header_2, ln=1, align='L')\n",
        "  pdf.ln(1)\n",
        "\n",
        "  all_packages = ''\n",
        "  for requirement in freeze(local_only=True):\n",
        "    all_packages = all_packages+requirement+', '\n",
        "  #print(all_packages)\n",
        "\n",
        "  #Main Packages\n",
        "  main_packages = ''\n",
        "  version_numbers = []\n",
        "  for name in ['tensorflow','numpy','keras']:\n",
        "    find_name=all_packages.find(name)\n",
        "    main_packages = main_packages+all_packages[find_name:all_packages.find(',',find_name)]+', '\n",
        "    #Version numbers only here:\n",
        "    version_numbers.append(all_packages[find_name+len(name)+2:all_packages.find(',',find_name)])\n",
        "\n",
        "  cuda_version = subprocess.run('nvcc --version',stdout=subprocess.PIPE, shell=True)\n",
        "  cuda_version = cuda_version.stdout.decode('utf-8')\n",
        "  cuda_version = cuda_version[cuda_version.find(', V')+3:-1]\n",
        "  gpu_name = subprocess.run('nvidia-smi',stdout=subprocess.PIPE, shell=True)\n",
        "  gpu_name = gpu_name.stdout.decode('utf-8')\n",
        "  gpu_name = gpu_name[gpu_name.find('Tesla'):gpu_name.find('Tesla')+10]\n",
        "  #print(cuda_version[cuda_version.find(', V')+3:-1])\n",
        "  #print(gpu_name)\n",
        "\n",
        "  if os.path.isdir(training_source):\n",
        "    shape = io.imread(training_source+'/'+os.listdir(training_source)[0]).shape\n",
        "  elif os.path.isfile(training_source):\n",
        "    shape = io.imread(training_source).shape\n",
        "  else:\n",
        "    print('Cannot read training data.')\n",
        "\n",
        "  dataset_size = len(train_generator)\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)+') with a batch size of '+str(batch_size)+' and a '+loss_function+' loss function, using the '+Network+' ZeroCostDL4Mic notebook (v '+Notebook_version[0]+') (von Chamier & Laine et al., 2020). Key python packages used include tensorflow (v '+version_numbers[0]+'), keras (v '+version_numbers[2]+'), numpy (v '+version_numbers[1]+'), cuda (v '+cuda_version+'). The training was accelerated using a '+gpu_name+'GPU.'\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)+') with a batch size of '+str(batch_size)+' and a '+loss_function+' loss function, using the '+Network+' ZeroCostDL4Mic notebook (v '+Notebook_version[0]+') (von Chamier & Laine et al., 2020). The model was retrained from a pretrained model. Key python packages used include tensorflow (v '+version_numbers[0]+'), keras (v '+version_numbers[2]+'), numpy (v '+version_numbers[1]+'), cuda (v '+cuda_version+'). The training was accelerated using a '+gpu_name+'GPU.'\n",
        "\n",
        "  pdf.set_font('')\n",
        "  pdf.set_font_size(10.)\n",
        "  pdf.multi_cell(190, 5, txt = text, align='L')\n",
        "  pdf.ln(1)\n",
        "  pdf.set_font('')\n",
        "  pdf.set_font('Arial', size = 10, style = 'B')\n",
        "  pdf.cell(28, 5, txt='Augmentation: ', ln=0)\n",
        "  pdf.set_font('')\n",
        "  if augmentation:\n",
        "    aug_text = 'The dataset was augmented by'\n",
        "    if add_gaussian_blur == True:\n",
        "      aug_text = aug_text+'\\n- gaussian blur'\n",
        "    if add_linear_contrast == True:\n",
        "      aug_text = aug_text+'\\n- linear contrast'\n",
        "    if add_additive_gaussian_noise == True:\n",
        "      aug_text = aug_text+'\\n- additive gaussian noise'\n",
        "    if augmenters != '':\n",
        "      aug_text = aug_text+'\\n- imgaug augmentations: '+augmenters\n",
        "    if add_elastic_deform == True:\n",
        "      aug_text = aug_text+'\\n- elastic deformation'\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=60% style=\"margin-left:0px;\">\n",
        "    <tr>\n",
        "      <th width = 50% align=\"left\">Parameter</th>\n",
        "      <th width = 50% align=\"left\">Value</th>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <td width = 50%>number_of_epochs</td>\n",
        "      <td width = 50%>{0}</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <td width = 50%>batch_size</td>\n",
        "      <td width = 50%>{1}</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <td width = 50%>patch_size</td>\n",
        "      <td width = 50%>{2}</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <td width = 50%>image_pre_processing</td>\n",
        "      <td width = 50%>{3}</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <td width = 50%>validation_split_in_percent</td>\n",
        "      <td width = 50%>{4}</td>\n",
        "    </tr>\n",
        "      <tr>\n",
        "      <td width = 50%>downscaling_in_xy</td>\n",
        "      <td width = 50%>{5}</td>\n",
        "    </tr>\n",
        "      <tr>\n",
        "      <td width = 50%>binary_target</td>\n",
        "      <td width = 50%>{6}</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <td width = 50%>loss_function</td>\n",
        "      <td width = 50%>{7}</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <td width = 50%>metrics</td>\n",
        "      <td width = 50%>{8}</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <td width = 50%>optimizer</td>\n",
        "      <td width = 50%>{9}</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <td width = 50%>checkpointing_period</td>\n",
        "      <td width = 50%>{10}</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <td width = 50%>save_best_only</td>\n",
        "      <td width = 50%>{11}</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <td width = 50%>resume_training</td>\n",
        "      <td width = 50%>{12}</td>\n",
        "    </tr>\n",
        "  </table>\n",
        "  \"\"\".format(number_of_epochs,batch_size,str(patch_size[0])+'x'+str(patch_size[1])+'x'+str(patch_size[2]),image_pre_processing, validation_split_in_percent, downscaling_in_xy, str(binary_target), loss_function, metrics, optimizer, checkpointing_period, str(save_best_only), str(resume_training))\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 (single slice)', ln=1)\n",
        "  pdf.ln(1)\n",
        "  exp_size = io.imread('/content/TrainingDataExample_Unet3D.png').shape\n",
        "  pdf.image('/content/TrainingDataExample_Unet3D.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 = '- Unet 3D: Çiçek, Özgün, et al. \"3D U-Net: learning dense volumetric segmentation from sparse annotation.\" International conference on medical image computing and computer-assisted intervention. Springer, Cham, 2016.'\n",
        "  pdf.multi_cell(190, 5, txt = ref_2, align='L')\n",
        "  # if Use_Data_augmentation:\n",
        "  #   ref_4 = '- Augmentor: Bloice, Marcus D., Christof Stocker, and Andreas Holzinger. \"Augmentor: an image augmentation library for machine learning.\" arXiv preprint arXiv:1708.04680 (2017).'\n",
        "  #   pdf.multi_cell(190, 5, txt = ref_4, align='L')\n",
        "  pdf.ln(3)\n",
        "  reminder = 'Important:\\nRemember to perform the quality control step on all newly trained models\\nPlease consider depositing your training dataset on Zenodo'\n",
        "  pdf.set_font('Arial', size = 11, style='B')\n",
        "  pdf.multi_cell(190, 5, txt=reminder, align='C')\n",
        "  pdf.ln(1)\n",
        "\n",
        "  pdf.output(model_path+'/'+model_name+'/'+model_name+'_training_report.pdf')\n",
        "\n",
        "  print('------------------------------')\n",
        "  print('PDF report exported in '+model_path+'/'+model_name+'/')\n",
        "\n",
        "\n",
        "def qc_pdf_export():\n",
        "  class MyFPDF(FPDF, HTMLMixin):\n",
        "    pass\n",
        "\n",
        "  pdf = MyFPDF()\n",
        "  pdf.add_page()\n",
        "  pdf.set_right_margin(-1)\n",
        "  pdf.set_font(\"Arial\", size = 11, style='B') \n",
        "\n",
        "  Network = 'U-Net 3D'\n",
        "\n",
        "  day = datetime.now()\n",
        "  datetime_str = str(day)[0:10]\n",
        "\n",
        "  Header = 'Quality Control report for '+Network+' model ('+qc_model_name+')\\nDate: '+datetime_str\n",
        "  pdf.multi_cell(180, 5, txt = Header, align = 'L') \n",
        "  pdf.ln(1)\n",
        "\n",
        "  all_packages = ''\n",
        "  for requirement in freeze(local_only=True):\n",
        "    all_packages = all_packages+requirement+', '\n",
        "\n",
        "  pdf.set_font('')\n",
        "  pdf.set_font('Arial', size = 11, style = 'B')\n",
        "  pdf.ln(2)\n",
        "  pdf.cell(190, 5, txt = 'Loss curves', ln=1, align='L')\n",
        "  pdf.ln(1)\n",
        "  if os.path.exists(os.path.join(qc_model_path,qc_model_name,'Quality Control')+'/lossCurvePlots.png'):\n",
        "    exp_size = io.imread(os.path.join(qc_model_path,qc_model_name,'Quality Control')+'/lossCurvePlots.png').shape\n",
        "    pdf.image(os.path.join(qc_model_path,qc_model_name,'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(os.path.join(qc_model_path,qc_model_name,'Quality Control')+'/QC_example_data.png').shape\n",
        "  pdf.image(os.path.join(qc_model_path,qc_model_name,'Quality Control')+'/QC_example_data.png', x = 16, y = None, w = round(exp_size[1]/8), h = round(exp_size[0]/8))\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 = 'IoU threshold optimisation', align='L', ln=1)\n",
        "  pdf.set_font('')\n",
        "  pdf.set_font_size(10.)\n",
        "  pdf.ln(1)\n",
        "  pdf.cell(120, 5, txt='Highest IoU is {:.4f} with a threshold of {}'.format(best_iou, best_thresh), align='L', ln=1)\n",
        "  pdf.ln(2)\n",
        "  exp_size = io.imread(os.path.join(qc_model_path,qc_model_name,'Quality Control')+'/QC_IoU_analysis.png').shape\n",
        "  pdf.image(os.path.join(qc_model_path,qc_model_name,'Quality Control')+'/QC_IoU_analysis.png', x=16, y=None, w = round(exp_size[1]/6), h = round(exp_size[0]/6))\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 = '- Unet 3D: Çiçek, Özgün, et al. \"3D U-Net: learning dense volumetric segmentation from sparse annotation.\" International conference on medical image computing and computer-assisted intervention. Springer, Cham, 2016.'\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(os.path.join(qc_model_path,qc_model_name,'Quality Control')+'/'+qc_model_name+'_QC_report.pdf')\n",
        "\n",
        "  print('------------------------------')\n",
        "  print('QC PDF report exported in '+os.path.join(qc_model_path,qc_model_name,'Quality Control')+'/')\n",
        "\n",
        "\n",
        "# -------------- Other definitions -----------\n",
        "W  = '\\033[0m'  # white (normal)\n",
        "R  = '\\033[31m' # red\n",
        "prediction_prefix = 'Predicted_'\n",
        "\n",
        "\n",
        "print('-------------------')\n",
        "print('U-Net 3D and dependencies installed.')\n",
        "\n",
        "# Colors for the warning messages\n",
        "class bcolors:\n",
        "  WARNING = '\\033[31m'\n",
        "  NORMAL = '\\033[0m'  # white (normal)\n",
        "  \n",
        "\n",
        "# Check if this is the latest version of the notebook\n",
        "# Latest_notebook_version = pd.read_csv(\"https://raw.githubusercontent.com/HenriquesLab/ZeroCostDL4Mic/master/Colab_notebooks/Latest_ZeroCostDL4Mic_Release.csv\")\n",
        "\n",
        "# if Notebook_version == list(Latest_notebook_version.columns):\n",
        "#   print(\"This notebook is up-to-date.\")\n",
        "\n",
        "# if not Notebook_version == list(Latest_notebook_version.columns):\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",
        "All_notebook_versions = pd.read_csv(\"https://raw.githubusercontent.com/HenriquesLab/ZeroCostDL4Mic/master/Colab_notebooks/Latest_Notebook_versions.csv\", dtype=str)\n",
        "print('Notebook version: '+Notebook_version)\n",
        "Latest_Notebook_version = All_notebook_versions[All_notebook_versions[\"Notebook\"] == Network]['Version'].iloc[0]\n",
        "print('Latest notebook version: '+Latest_Notebook_version)\n",
        "if Notebook_version == Latest_Notebook_version:\n",
        "  print(\"This notebook is up-to-date.\")\n",
        "else:\n",
        "  print(bcolors.WARNING +\"A new version of this notebook has been released. We recommend that you download it at https://github.com/HenriquesLab/ZeroCostDL4Mic/wiki\")\n",
        "\n",
        "\n",
        "\n",
        "# Build requirements file for local run\n",
        "after = [str(m) for m in sys.modules]\n",
        "build_requirements_file(before, after)\n",
        "!pip3 freeze > requirements.txt"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "n4yWFoJNnoin"
      },
      "source": [
        "# **2. Complete 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",
        "\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\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sNIVx8_CLolt"
      },
      "source": [
        "## **2.2. Mount Google Drive**\n",
        "---\n",
        "<font size = 4> To use this notebook with your **own data**, place it in a folder on **Google Drive** following one of the directory structures outlined in **Section 0**.\n",
        "\n",
        "1. <font size = 4> **Run** the **cell** below to mount your Google Drive and follow the link. \n",
        "\n",
        "2. <font size = 4>**Sign in** to your Google account and press 'Allow'. \n",
        "\n",
        "3. <font size = 4>Next, copy the **authorization code**, paste it into the cell and press enter. This will allow Colab to read and write data from and to your Google Drive. \n",
        "\n",
        "4. <font size = 4> Once this is done, your data can be viewed in the **Files tab** on the top left of the notebook after hitting 'Refresh'."
      ]
    },
    {
      "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",
        "# mount user's Google Drive to Google Colab.\n",
        "from google.colab import drive\n",
        "drive.mount('/content/gdrive')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HLYcZR9gMv42"
      },
      "source": [
        "# **3. Select your parameters and paths**\n",
        "---"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AuESFimvMv43"
      },
      "source": [
        "## **3.1. Choosing parameters**\n",
        "\n",
        "---\n",
        "\n",
        "### **Paths to training data and model**\n",
        "\n",
        "* <font size = 4>**`training_source`** and **`training_target`** specify the paths to the training data. They can either be a single multipage TIFF file each or directories containing various multipage TIFF files in which case target and source files must be named identically within the respective directories. See Section 0 for a detailed description of the necessary directory structure.\n",
        "\n",
        "* <font size = 4>**`model_name`** will be used when naming checkpoints. Adhere to a `lower_case_with_underscores` naming convention and beware of using the name of an existing model within the same folder, as it will be overwritten.\n",
        "\n",
        "* <font size = 4>**`model_path`** specifies the directory where the model checkpoints and quality control logs will be saved.\n",
        "\n",
        "\n",
        "<font size = 4>**Note:** You can copy paths from the 'Files' tab by right-clicking any folder or file and selecting 'Copy path'. \n",
        "\n",
        "### **Training parameters**\n",
        "\n",
        "* <font size = 4>**`number_of_epochs`** is the number of times the entire training data will be seen by the model. *Default: >100*\n",
        "\n",
        "* <font size = 4>**`batch_size`** is the number of training patches of size `patch_size` that will be bundled together at each training step. *Default: 1*\n",
        "\n",
        "* <font size = 4>**`patch_size`** specifies the size of the three-dimensional training patches in (x, y, z) that will be fed to the model. In order to avoid errors, preferably use a square aspect ratio or stick to the advanced parameters. *Default: <(512, 512, 16)*\n",
        "\n",
        "* <font size = 4>**`validation_split_in_percent`** is the relative amount of training data that will be set aside for validation. *Default: 20* \n",
        "\n",
        "* <font size = 4>**`downscaling_in_xy`** downscales the training images by the specified amount in x and y. This is useful to enforce isotropic pixel-size if the z resolution is lower than the xy resolution in the training volume or to capture a larger field-of-view while decreasing the memory requirements. *Default: 1*\n",
        "\n",
        "* <font size = 4>**`image_pre_processing`** selects whether the training images are randomly cropped during training or resized to `patch_size`. Choose `randomly crop to patch_size` to shrink the field-of-view of the training images to the `patch_size`. *Default: resize to patch_size* \n",
        "\n",
        "* <font size = 4>**`binary_target`** forces the target image to be binary. Choose this if your model is trained to perform binary segmentation tasks *Default: True* \n",
        "\n",
        "* <font size = 4>**`loss_function`** defines the loss. Read more [here](https://keras.io/api/losses/). *Default: weighted_binary_crossentropy* \n",
        "\n",
        "* <font size = 4>**`metrics`** defines the metric. Read more [here](https://keras.io/api/metrics/). *Default: dice* \n",
        "\n",
        "* <font size = 4>**`optimizer`** defines the optimizer. Read more [here](https://keras.io/api/optimizers/). *Default: adam* \n",
        "\n",
        "<font size = 4>**Note:** If a *ResourceExhaustedError* is raised in Section 4.1. during training, decrease `batch_size` and `patch_size`. Decrease `batch_size` first and if the error persists at `batch_size = 1`, reduce the `patch_size`.  \n",
        "\n",
        "<font size = 4>**Note:** The number of steps per epoch are calculated as `floor(augment_factor * (1 - validation_split) * num_of_slices / batch_size)` if `image_pre_processing` is `resize to patch_size` where `augment_factor` is three if `apply_data_augmentation` is `True` and one otherwise. The `num_of_slices` is the overall number of slices (z-depth) in the training set across all provided image volumes. If `image_pre_processing` is `randomly crop to patch_size`, the number of steps per epoch are calculated as `floor(augment_factor * volume / (crop_volume * batch_size))` where `volume` is the overall volume of the training data in pixels accounting for the validation split and `crop_volume` is defined as the volume in pixels based on the specified `patch_size`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "ewpNJ_I0Mv47"
      },
      "outputs": [],
      "source": [
        "#@markdown ###Path to training data:\n",
        "training_source = \"\" #@param {type:\"string\"}\n",
        "training_target = \"\" #@param {type:\"string\"}\n",
        "\n",
        "#@markdown ---\n",
        "\n",
        "#@markdown ###Model name and path to model folder:\n",
        "model_name = \"\" #@param {type:\"string\"}\n",
        "model_path = \"\" #@param {type:\"string\"}\n",
        "\n",
        "full_model_path = os.path.join(model_path, model_name)\n",
        "\n",
        "#@markdown ---\n",
        "\n",
        "#@markdown ###Training parameters\n",
        "number_of_epochs =   2#@param {type:\"number\"}\n",
        "\n",
        "#@markdown ###Default advanced parameters\n",
        "use_default_advanced_parameters = False #@param {type:\"boolean\"}\n",
        "\n",
        "#@markdown <font size = 3>If not, please change:\n",
        "\n",
        "batch_size =  1#@param {type:\"number\"}\n",
        "patch_size = (512,512,8) #@param {type:\"number\"} # in pixels\n",
        "training_shape = patch_size + (1,)\n",
        "image_pre_processing = 'resize to patch_size' #@param [\"randomly crop to patch_size\", \"resize to patch_size\"]\n",
        "\n",
        "validation_split_in_percent = 20 #@param{type:\"number\"}\n",
        "downscaling_in_xy =  1#@param {type:\"number\"} # in pixels\n",
        "\n",
        "binary_target = True #@param {type:\"boolean\"}\n",
        "\n",
        "loss_function = 'weighted_binary_crossentropy' #@param [\"weighted_binary_crossentropy\", \"binary_crossentropy\", \"categorical_crossentropy\", \"sparse_categorical_crossentropy\", \"mean_squared_error\", \"mean_absolute_error\"]\n",
        "\n",
        "metrics = 'dice' #@param [\"dice\", \"accuracy\"]\n",
        "\n",
        "optimizer = 'adam' #@param [\"adam\", \"sgd\", \"rmsprop\"]\n",
        "\n",
        "learning_rate = 0.001 #@param{type:\"number\"}\n",
        "\n",
        "if image_pre_processing == \"randomly crop to patch_size\":\n",
        "    random_crop = True\n",
        "else:\n",
        "    random_crop = False\n",
        "\n",
        "if use_default_advanced_parameters: \n",
        "    print(\"Default advanced parameters enabled\")\n",
        "    batch_size = 3\n",
        "    training_shape = (256,256,8,1)\n",
        "    validation_split_in_percent = 20\n",
        "    downscaling_in_xy = 1\n",
        "    random_crop = True\n",
        "    binary_target = True\n",
        "    loss_function = 'weighted_binary_crossentropy'\n",
        "    metrics = 'dice'\n",
        "    optimizer = 'adam'\n",
        "    learning_rate = 0.001 \n",
        "#@markdown ###Checkpointing parameters\n",
        "#checkpointing_period = 1 #@param {type:\"number\"}\n",
        "checkpointing_period = \"epoch\"\n",
        "#@markdown  <font size = 3>If chosen, only the best checkpoint is saved. Otherwise a checkpoint is saved every epoch:\n",
        "save_best_only = False #@param {type:\"boolean\"}\n",
        "\n",
        "#@markdown ###Resume training\n",
        "#@markdown <font size = 3>Choose if training was interrupted:\n",
        "resume_training = False #@param {type:\"boolean\"}\n",
        "\n",
        "#@markdown ###Transfer learning\n",
        "#@markdown <font size = 3>For transfer learning, do not select resume_training and specify a checkpoint_path below. \n",
        "\n",
        "#@markdown <font size = 3> - If the model is already downloaded or is locally available, please specify the path to the .h5 file. \n",
        "\n",
        "#@markdown <font size = 3> - To use a model from the BioImage Model Zoo, write the model ID. For example: 10.5281/zenodo.5749843\n",
        "\n",
        "pretrained_model_choice = \"Model_from_file\" #@param [\"Model_from_file\", \"bioimageio_model\"]\n",
        "checkpoint_path = \"\" #@param {type:\"string\"}\n",
        "model_id = \"\" #@param {type:\"string\"}\n",
        "# --------------------- Load the model from a bioimageio model (can be path on drive or url / doi) ---\n",
        "if pretrained_model_choice == \"bioimageio_model\":\n",
        "  from bioimageio.core import load_raw_resource_description\n",
        "  from zipfile import ZipFile\n",
        "  import requests\n",
        "\n",
        "  model_spec = load_raw_resource_description(model_id)\n",
        "  if \"keras_hdf5\" not in model_spec.weights:\n",
        "    print(\"Invalid bioimageio model\")\n",
        "  else:\n",
        "    url = model_spec.weights[\"keras_hdf5\"].source\n",
        "    r = requests.get(url, allow_redirects=True)\n",
        "    open(\"keras_model.h5\", 'wb').write(r.content)\n",
        "    checkpoint_path = \"keras_model.h5\"\n",
        "\n",
        "if resume_training and checkpoint_path != \"\":\n",
        "    print('If resume_training is True while checkpoint_path is specified, resume_training will be set to False!')\n",
        "    resume_training = False\n",
        "\n",
        "# Retrieve last checkpoint\n",
        "if resume_training:\n",
        "    try:\n",
        "      ckpt_dir_list = glob(full_model_path + '/ckpt/*')\n",
        "      ckpt_dir_list.sort()\n",
        "      last_ckpt_path = ckpt_dir_list[-1]\n",
        "      print('Training will resume from checkpoint:', os.path.basename(last_ckpt_path))\n",
        "    except IndexError:\n",
        "      last_ckpt_path=None\n",
        "      print('CheckpointError: No previous checkpoints were found, training from scratch.')\n",
        "elif not resume_training and checkpoint_path != \"\":\n",
        "    last_ckpt_path = checkpoint_path\n",
        "    assert os.path.isfile(last_ckpt_path), 'checkpoint_path does not exist!'\n",
        "else:\n",
        "    last_ckpt_path=None\n",
        "\n",
        "# Instantiate Unet3D \n",
        "model = Unet3D(shape=training_shape)\n",
        "\n",
        "#here we check that no model with the same name already exist\n",
        "if not resume_training and os.path.exists(full_model_path): \n",
        "    print(bcolors.WARNING+'The model folder already exists and will be overwritten.'+bcolors.NORMAL)\n",
        "    # print('!! WARNING: Folder already exists and will be overwritten !!') \n",
        "    # shutil.rmtree(full_model_path)\n",
        "\n",
        "# if not os.path.exists(full_model_path):\n",
        "#     os.makedirs(full_model_path)\n",
        "\n",
        "# Show sample image\n",
        "if os.path.isdir(training_source):\n",
        "    training_source_sample = sorted(glob(os.path.join(training_source, '*')))[0]\n",
        "    training_target_sample = sorted(glob(os.path.join(training_target, '*')))[0]\n",
        "else:\n",
        "    training_source_sample = training_source\n",
        "    training_target_sample = training_target\n",
        "\n",
        "src_sample = tifffile.imread(training_source_sample)\n",
        "src_sample = model._min_max_scaling(src_sample)\n",
        "if binary_target:\n",
        "    tgt_sample = tifffile.imread(training_target_sample).astype(np.bool)\n",
        "else:\n",
        "    tgt_sample = tifffile.imread(training_target_sample)\n",
        "\n",
        "src_down = transform.downscale_local_mean(src_sample[0], (downscaling_in_xy, downscaling_in_xy))\n",
        "tgt_down = transform.downscale_local_mean(tgt_sample[0], (downscaling_in_xy, downscaling_in_xy))   \n",
        "\n",
        "if random_crop:\n",
        "    true_patch_size = None\n",
        "\n",
        "    if src_down.shape[0] == training_shape[0]:\n",
        "      x_rand = 0\n",
        "    if src_down.shape[1] == training_shape[1]:\n",
        "      y_rand = 0\n",
        "    if src_down.shape[0] > training_shape[0]:\n",
        "      x_rand = np.random.randint(src_down.shape[0] - training_shape[0])\n",
        "    if src_down.shape[1] > training_shape[1]:\n",
        "      y_rand = np.random.randint(src_down.shape[1] - training_shape[1])\n",
        "    if src_down.shape[0] < training_shape[0] or src_down.shape[1] < training_shape[1]:\n",
        "      raise ValueError('Patch shape larger than (downscaled) source shape')\n",
        "else:\n",
        "    true_patch_size = src_down.shape\n",
        "\n",
        "def scroll_in_z(z):\n",
        "    src_down = transform.downscale_local_mean(src_sample[z-1], (downscaling_in_xy,downscaling_in_xy))\n",
        "    tgt_down = transform.downscale_local_mean(tgt_sample[z-1], (downscaling_in_xy,downscaling_in_xy))       \n",
        "    if random_crop:\n",
        "        src_slice = src_down[x_rand:training_shape[0]+x_rand, y_rand:training_shape[1]+y_rand]\n",
        "        tgt_slice = tgt_down[x_rand:training_shape[0]+x_rand, y_rand:training_shape[1]+y_rand]\n",
        "    else:\n",
        "        \n",
        "        src_slice = transform.resize(src_down, (training_shape[0], training_shape[1]), mode='constant', preserve_range=True)\n",
        "        tgt_slice = transform.resize(tgt_down, (training_shape[0], training_shape[1]), mode='constant', preserve_range=True)\n",
        "\n",
        "    f=plt.figure(figsize=(16,8))\n",
        "    plt.subplot(1,2,1)\n",
        "    plt.imshow(src_slice, cmap='gray')\n",
        "    plt.title('Training source (z = ' + str(z) + ')', fontsize=15)\n",
        "    plt.axis('off')\n",
        "\n",
        "    plt.subplot(1,2,2)\n",
        "    plt.imshow(tgt_slice, cmap='magma')\n",
        "    plt.title('Training target (z = ' + str(z) + ')', fontsize=15)\n",
        "    plt.axis('off')\n",
        "    plt.savefig('/content/TrainingDataExample_Unet3D.png',bbox_inches='tight',pad_inches=0)\n",
        "    #plt.close()\n",
        "\n",
        "print('This is what the training images will look like with the chosen settings')\n",
        "interact(scroll_in_z, z=widgets.IntSlider(min=1, max=src_sample.shape[0], step=1, value=0));\n",
        "plt.show()\n",
        "#Create a copy of an example slice and close the display.\n",
        "scroll_in_z(z=int(src_sample.shape[0]/2))\n",
        "# If you close the display, then the users can't interactively inspect the data\n",
        "# plt.close()\n",
        "\n",
        "# Save model parameters\n",
        "params =  {'training_source': training_source,\n",
        "           'training_target': training_target,\n",
        "           'model_name': model_name,\n",
        "           'model_path': model_path,\n",
        "           'number_of_epochs': number_of_epochs,\n",
        "           'batch_size': batch_size,\n",
        "           'training_shape': training_shape,\n",
        "           'downscaling': downscaling_in_xy,\n",
        "           'true_patch_size': true_patch_size,\n",
        "           'val_split': validation_split_in_percent/100,\n",
        "           'random_crop': random_crop}\n",
        "\n",
        "params_df = pd.DataFrame.from_dict(params, orient='index')\n",
        "\n",
        "# apply_data_augmentation = False\n",
        "# pdf_export(augmentation = apply_data_augmentation, pretrained_model = resume_training)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "w_jCy7xOx2g3"
      },
      "source": [
        "## **3.2. Data augmentation**\n",
        " \n",
        "---\n",
        "<font size = 4> Augmenting the training data increases robustness of the model by simulating possible variations within the training data which avoids it from overfitting on small datasets. We therefore strongly recommended augmenting the data and making sure that the applied augmentations are reasonable.\n",
        "\n",
        "* <font size = 4>**Gaussian blur** blurs images using Gaussian kernels with a sigma of `gaussian_sigma`. This augmentation step is applied with a probability of `gaussian_frequency`. Read more [here](https://imgaug.readthedocs.io/en/latest/source/overview/blur.html#gaussianblur).\n",
        "\n",
        "* <font size = 4>**Linear contrast** modifies the contrast of images according to `127 + alpha *(pixel_value-127)`, where `pixel_value` and `alpha` are sampled uniformly from the interval `[contrast_min, contrast_max]`. This augmentation step is applied with a probability of `contrast_frequency`. Read more [here](https://imgaug.readthedocs.io/en/latest/source/overview/contrast.html#linearcontrast).\n",
        "\n",
        "* <font size = 4>**Additive Gaussian noise** adds Gaussian noise sampled once per pixel from a normal distribution `N(0, s)`, where `s` is sampled from `[scale_min, scale_max]`. This augmentation step is applied with a probability of `noise_frequency`. Read more [here](https://imgaug.readthedocs.io/en/latest/source/overview/arithmetic.html#additivegaussiannoise).\n",
        "\n",
        "* <font size = 4>**Add custom augmenters** allows you to create a custom augmentation pipeline using the [augmenters available in the imagug library](https://imgaug.readthedocs.io/en/latest/source/overview_of_augmenters.html).\n",
        "In the example above, the augmentation pipeline is equivalent to: \n",
        "```\n",
        "seq = iaa.Sequential([\n",
        "    iaa.Sometimes(0.3, iaa.GammaContrast((0.5, 2.0)), \n",
        "    iaa.Sometimes(0.4, iaa.AverageBlur((0.5, 2.0)), \n",
        "    iaa.Sometimes(0.5, iaa.LinearContrast((0.4, 1.6)), \n",
        "], random_order=True)\n",
        "```\n",
        "<font size = 4> Note that there is no limit on the number of augmenters that can be chained together and that individual augmenter and parameter entries must be separated by `;`. Custom augmenters do not overwrite the preset augmentation steps (*Gaussian blur*, *Linear contrast* or *Additive Gaussian noise*). Also, the augmenters, augmenter parameters and augmenter frequencies must be entered such that each position within the string corresponds to the same augmentation step.\n",
        "\n",
        "* <font size = 4>**`apply_data_augmentation`** ensures that data augmentation is randomly applied to the training data at each training step. This includes inverting the order of the slices within a training patch, as well as applying any augmenters that are added. *Default: True*\n",
        "\n",
        "* <font size = 4>**`add_elastic_deform`** ensures that elastic grid-based deformations are applied as described in the original 3D U-Net paper. *Default: True*"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "DMqWq5-AxnFU"
      },
      "outputs": [],
      "source": [
        "#@markdown ##**Augmentation options**\n",
        "\n",
        "#@markdown ###Data augmentation\n",
        "\n",
        "apply_data_augmentation = True #@param {type:\"boolean\"}\n",
        "\n",
        "# List of augmentations\n",
        "augmentations = []\n",
        "\n",
        "#@markdown ###Gaussian blur\n",
        "add_gaussian_blur = True #@param {type:\"boolean\"}\n",
        "gaussian_sigma =   0.7#@param {type:\"number\"}\n",
        "gaussian_frequency = 0.5 #@param {type:\"number\"}\n",
        "\n",
        "if add_gaussian_blur:\n",
        "    augmentations.append(iaa.Sometimes(gaussian_frequency, iaa.GaussianBlur(sigma=(0, gaussian_sigma))))\n",
        "\n",
        "#@markdown ###Linear contrast\n",
        "add_linear_contrast = True #@param {type:\"boolean\"}\n",
        "contrast_min =  0.4 #@param {type:\"number\"}\n",
        "contrast_max =   1.6#@param {type:\"number\"}\n",
        "contrast_frequency = 0.5 #@param {type:\"number\"}\n",
        "\n",
        "if add_linear_contrast:\n",
        "    augmentations.append(iaa.Sometimes(contrast_frequency, iaa.LinearContrast((contrast_min, contrast_max))))\n",
        "\n",
        "#@markdown ###Additive Gaussian noise\n",
        "add_additive_gaussian_noise = True #@param {type:\"boolean\"}\n",
        "scale_min =  0 #@param {type:\"number\"}\n",
        "scale_max =  0.05 #@param {type:\"number\"}\n",
        "noise_frequency = 0.5 #@param {type:\"number\"}\n",
        "\n",
        "if add_additive_gaussian_noise:\n",
        "    augmentations.append(iaa.Sometimes(noise_frequency, iaa.AdditiveGaussianNoise(scale=(scale_min, scale_max))))\n",
        "\n",
        "#@markdown ###Add custom augmenters\n",
        "add_custom_augmenters = False #@param {type:\"boolean\"}    \n",
        "augmenters = \"\" #@param {type:\"string\"}\n",
        "\n",
        "if add_custom_augmenters:\n",
        "\n",
        "    augmenter_params = \"\" #@param {type:\"string\"}\n",
        "\n",
        "    augmenter_frequency = \"\" #@param {type:\"string\"}\n",
        "\n",
        "    aug_lst = augmenters.split(';')\n",
        "    aug_params_lst = augmenter_params.split(';')\n",
        "    aug_freq_lst = augmenter_frequency.split(';')\n",
        "\n",
        "    assert len(aug_lst) == len(aug_params_lst) and len(aug_lst) == len(aug_freq_lst), 'The number of arguments in augmenters, augmenter_params and augmenter_frequency are not the same!'\n",
        "\n",
        "    for __, (aug, param, freq) in enumerate(zip(aug_lst, aug_params_lst, aug_freq_lst)):\n",
        "        aug, param, freq = aug.strip(), param.strip(), freq.strip() \n",
        "        aug_func = iaa.Sometimes(eval(freq), getattr(iaa, aug)(eval(param)))\n",
        "        augmentations.append(aug_func)\n",
        "\n",
        "#@markdown ###Elastic deformations\n",
        "add_elastic_deform = False #@param {type:\"boolean\"}\n",
        "sigma =  2#@param {type:\"number\"}\n",
        "points =  2#@param {type:\"number\"}\n",
        "order =  1#@param {type:\"number\"}\n",
        "\n",
        "if add_elastic_deform:\n",
        "    deform_params = (sigma, points, order)\n",
        "else:\n",
        "    deform_params = None\n",
        "\n",
        "train_generator = MultiPageTiffGenerator(training_source,\n",
        "                                         training_target,\n",
        "                                         batch_size=batch_size,\n",
        "                                         shape=training_shape,\n",
        "                                         augment=apply_data_augmentation,\n",
        "                                         augmentations=augmentations,\n",
        "                                         deform_augment=add_elastic_deform,\n",
        "                                         deform_augmentation_params=deform_params,\n",
        "                                         val_split=validation_split_in_percent/100,\n",
        "                                         random_crop=random_crop,\n",
        "                                         downscale=downscaling_in_xy,\n",
        "                                         binary_target=binary_target)\n",
        "\n",
        "val_generator = MultiPageTiffGenerator(training_source,\n",
        "                                       training_target,\n",
        "                                       batch_size=batch_size,\n",
        "                                       shape=training_shape,\n",
        "                                       val_split=validation_split_in_percent/100,\n",
        "                                       is_val=True,\n",
        "                                       random_crop=random_crop,\n",
        "                                       downscale=downscaling_in_xy,\n",
        "                                       binary_target=binary_target)\n",
        "\n",
        "\n",
        "if apply_data_augmentation:\n",
        "  print('Data augmentation enabled.')\n",
        "  sample_src_aug, sample_tgt_aug = train_generator.sample_augmentation(random.randint(0, len(train_generator)))\n",
        "\n",
        "  def scroll_in_z(z):\n",
        "      f=plt.figure(figsize=(16,8))\n",
        "      plt.subplot(1,2,1)\n",
        "      plt.imshow(sample_src_aug[0,:,:,z-1,0], cmap='gray')\n",
        "      plt.title('Sample augmented source (z = ' + str(z) + ')', fontsize=15)\n",
        "      plt.axis('off')\n",
        "\n",
        "      plt.subplot(1,2,2)\n",
        "      plt.imshow(sample_tgt_aug[0,:,:,z-1,0], cmap='magma')\n",
        "      plt.title('Sample training target (z = ' + str(z) + ')', fontsize=15)\n",
        "      plt.axis('off')\n",
        "\n",
        "  print('This is what the augmented training images will look like with the chosen settings')\n",
        "  interact(scroll_in_z, z=widgets.IntSlider(min=1, max=sample_src_aug.shape[3], step=1, value=0));\n",
        "\n",
        "else:\n",
        "  print('Data augmentation disabled.')\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MCGklf1vZf2M"
      },
      "source": [
        "# **4. Train the network**\n",
        "---\n",
        "\n",
        "<font size = 4>**CRITICAL NOTE:** Google Colab has a time limit for processing (to prevent using GPU power for datamining). Training times must be less than 12 hours! If training takes longer than 12 hours, please decrease `number_of_epochs`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1KYOuygETJkT"
      },
      "source": [
        "## **4.1. Show model and start training**\n",
        "---\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "lIUAOJ_LMv5E"
      },
      "outputs": [],
      "source": [
        "#@markdown ## Show model summary\n",
        "model.summary()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "CyQI4ssarUp4"
      },
      "outputs": [],
      "source": [
        "#@markdown ##Start training\n",
        "\n",
        "#here we check that no model with the same name already exist, if so delete\n",
        "if not resume_training and os.path.exists(full_model_path): \n",
        "    shutil.rmtree(full_model_path)\n",
        "    print(bcolors.WARNING+'!! WARNING: Folder already exists and has been overwritten !!'+bcolors.NORMAL) \n",
        "\n",
        "if not os.path.exists(full_model_path):\n",
        "    os.makedirs(full_model_path)\n",
        "\n",
        "pdf_export(augmentation = apply_data_augmentation, pretrained_model = resume_training)\n",
        "\n",
        "# Save file\n",
        "params_df.to_csv(os.path.join(full_model_path, 'params.csv'))\n",
        "\n",
        "start = time.time()\n",
        "# Start Training\n",
        "model.train(epochs=number_of_epochs,\n",
        "            batch_size=batch_size,\n",
        "            train_generator=train_generator,\n",
        "            val_generator=val_generator,\n",
        "            model_path=model_path,\n",
        "            model_name=model_name,\n",
        "            loss=loss_function,\n",
        "            metrics=metrics,\n",
        "            optimizer=optimizer,\n",
        "            learning_rate=learning_rate,\n",
        "            ckpt_period=checkpointing_period,\n",
        "            save_best_ckpt_only=save_best_only,\n",
        "            ckpt_path=last_ckpt_path)\n",
        "\n",
        "print('Training successfully completed!')\n",
        "dt = time.time() - start\n",
        "mins, sec = divmod(dt, 60) \n",
        "hour, mins = divmod(mins, 60) \n",
        "print(\"Time elapsed:\",hour, \"hour(s)\",mins,\"min(s)\",round(sec),\"sec(s)\")\n",
        "\n",
        "#Create a pdf document with training summary\n",
        "\n",
        "pdf_export(trained = True, augmentation = apply_data_augmentation, pretrained_model = resume_training)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0Dfn8ZsEMv5d"
      },
      "source": [
        "##**4.2. Download your model from Google Drive**\n",
        "\n",
        "---\n",
        "<font size = 4>Once training is complete, the trained model is automatically saved to your Google Drive, in the **`model_path`** folder that was specified in Section 3. Download the folder to avoid any unwanted surprises, since the data can be erased if you train another model using the same `model_path`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "iwNmp1PUzRDQ",
        "scrolled": true
      },
      "outputs": [],
      "source": [
        "#@markdown ##Download model directory\n",
        "#@markdown 1.  <font size = 4>Specify the model_path in `model_path_download` otherwise the model sepcified in Section 3.1 will be downloaded\n",
        "#@markdown 2.  <font size = 4>Run this cell to zip the model directory\n",
        "#@markdown 3.  <font size = 4>Download the zipped file from the *Files* tab on the left\n",
        "\n",
        "from google.colab import files\n",
        "\n",
        "model_path_download = \"\" #@param {type:\"string\"}\n",
        "\n",
        "if len(model_path_download) == 0:\n",
        "    model_path_download = full_model_path\n",
        "\n",
        "model_name_download = os.path.basename(model_path_download)\n",
        "\n",
        "print('Zipping', model_name_download)\n",
        "\n",
        "zip_model_path = model_name_download + '.zip'\n",
        "\n",
        "!zip -r \"$zip_model_path\" \"$model_path_download\"\n",
        "\n",
        "print('Successfully saved zipped model directory as', zip_model_path)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_0Hynw3-xHp1"
      },
      "source": [
        "# **5. Evaluate your model**\n",
        "---\n",
        "\n",
        "<font size = 4>In this section the newly trained model can be assessed for performance. This involves inspecting the loss function in Section 5.1. and employing more advanced metrics in Section 5.2.\n",
        "\n",
        "<font size = 4>**We highly recommend performing quality control on all newly trained models.**\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "eAJzMwPA6tlH"
      },
      "outputs": [],
      "source": [
        "#@markdown ###Model to be evaluated:\n",
        "#@markdown <font size = 3>If left blank, the latest model defined in Section 3 will be evaluated:\n",
        "\n",
        "qc_model_name = \"\" #@param {type:\"string\"}\n",
        "qc_model_path = \"\" #@param {type:\"string\"}\n",
        "\n",
        "if len(qc_model_path) == 0 and len(qc_model_name) == 0:\n",
        "    qc_model_name = model_name\n",
        "    qc_model_path = model_path\n",
        "\n",
        "full_qc_model_path = os.path.join(qc_model_path, qc_model_name)\n",
        "\n",
        "if os.path.exists(full_qc_model_path):\n",
        "    print(qc_model_name + ' will be evaluated')\n",
        "else:\n",
        "    W  = '\\033[0m'  # white (normal)\n",
        "    R  = '\\033[31m' # red\n",
        "    print(R+'!! WARNING: The chosen model does not exist !!'+W)\n",
        "    print('Please make sure you provide a valid model path and model name before proceeding further.')\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dhJROwlAMv5o"
      },
      "source": [
        "## **5.1. Inspecting loss function**\n",
        "---\n",
        "\n",
        "<font size = 4>**The training loss** is the error between prediction and target after each epoch calculated across the training data while the **validation loss** calculates the error on the (unseen) validation data. During training these values should decrease until converging at which point the model has been sufficiently trained. If the validation loss starts increasing while the training loss has plateaued, the model has overfit on the training data which reduces its ability to generalise. Aim to halt training before this point.\n",
        "\n",
        "<font size = 4>**Note:** For a more in-depth explanation please refer to [this review](https://www.ncbi.nlm.nih.gov/pmc/articles/PMC6381354/) by Nichols et al.\n",
        "\n",
        "\n",
        "<font size = 4>The accuracy is another performance metric that is calculated after each epoch. We use the [Sørensen–Dice coefficient](https://en.wikipedia.org/wiki/S%C3%B8rensen%E2%80%93Dice_coefficient) to score the prediction accuracy. \n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "vMzSP50kMv5p"
      },
      "outputs": [],
      "source": [
        "#@markdown ##Visualise loss and accuracy\n",
        "lossDataFromCSV = []\n",
        "vallossDataFromCSV = []\n",
        "accuracyDataFromCSV = []\n",
        "valaccuracyDataFromCSV = []\n",
        "\n",
        "with open(full_qc_model_path + '/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[2]))\n",
        "        vallossDataFromCSV.append(float(row[4]))\n",
        "        accuracyDataFromCSV.append(float(row[1]))\n",
        "        valaccuracyDataFromCSV.append(float(row[3]))\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 and validation loss', fontsize=14)\n",
        "plt.ylabel('Loss', fontsize=12)\n",
        "plt.xlabel('Epochs', fontsize=12)\n",
        "plt.legend()\n",
        "\n",
        "plt.subplot(2,1,2)\n",
        "plt.plot(epochNumber,accuracyDataFromCSV, label='Training accuracy')\n",
        "plt.plot(epochNumber,valaccuracyDataFromCSV, label='Validation accuracy')\n",
        "plt.title('Training and validation accuracy', fontsize=14)\n",
        "plt.ylabel('Dice', fontsize=12)\n",
        "plt.xlabel('Epochs', fontsize=12)\n",
        "plt.legend()\n",
        "plt.savefig(full_qc_model_path + '/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 provide both a visual indication of the model performance by comparing the overlay of the predicted and source volume."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "w90MdriMxhjD"
      },
      "outputs": [],
      "source": [
        "#@markdown ##Compare prediction and ground-truth on testing data\n",
        "\n",
        "#@markdown <font size = 4>Provide an unseen annotated dataset to determine the performance of the model:\n",
        "\n",
        "testing_source = \"\" #@param{type:\"string\"}\n",
        "testing_target = \"\" #@param{type:\"string\"}\n",
        "\n",
        "qc_dir = full_qc_model_path + '/Quality Control'\n",
        "predict_dir = qc_dir + '/Prediction'\n",
        "if os.path.exists(predict_dir):\n",
        "    shutil.rmtree(predict_dir)\n",
        "\n",
        "os.makedirs(predict_dir)\n",
        "\n",
        "# predict_dir + '/' + \n",
        "predict_path = os.path.splitext(os.path.basename(testing_source))[0] + '_prediction.tif'\n",
        "\n",
        "def last_chars(x):\n",
        "    return(x[-11:])\n",
        "\n",
        "try:\n",
        "    ckpt_dir_list = glob(full_qc_model_path + '/ckpt/*')\n",
        "    ckpt_dir_list.sort(key=last_chars)\n",
        "    last_ckpt_path = ckpt_dir_list[0]\n",
        "    print('Predicting from checkpoint:', os.path.basename(last_ckpt_path))\n",
        "except IndexError:\n",
        "    raise CheckpointError('No previous checkpoints were found, please retrain model.')\n",
        "\n",
        "# Load parameters\n",
        "params = pd.read_csv(os.path.join(full_qc_model_path, 'params.csv'), names=['val'], header=0, index_col=0)   \n",
        "\n",
        "model = Unet3D(shape=params.loc['training_shape', 'val'])\n",
        "\n",
        "prediction = model.predict(testing_source, last_ckpt_path, \n",
        "                           downscaling=params.loc['downscaling', 'val'],\n",
        "                           true_patch_size=params.loc['true_patch_size', 'val'])\n",
        "\n",
        "tifffile.imwrite(predict_path, prediction.astype('float32'), imagej=True)\n",
        "\n",
        "print('Predicted images!')\n",
        "\n",
        "qc_metrics_path = full_qc_model_path + '/Quality Control/QC_metrics_' + qc_model_name + '.csv'\n",
        "\n",
        "test_target = tifffile.imread(testing_target)\n",
        "test_source = tifffile.imread(testing_source)\n",
        "test_prediction = tifffile.imread(predict_path)\n",
        "\n",
        "def scroll_in_z(z):\n",
        "\n",
        "    plt.figure(figsize=(25,5))\n",
        "    # Source\n",
        "    plt.subplot(1,4,1)\n",
        "    plt.axis('off')\n",
        "    plt.imshow(test_source[z-1], cmap='gray')\n",
        "    plt.title('Source (z = ' + str(z) + ')', fontsize=15)\n",
        "\n",
        "    # Target (Ground-truth)\n",
        "    plt.subplot(1,4,2)\n",
        "    plt.axis('off')\n",
        "    plt.imshow(test_target[z-1], cmap='magma')\n",
        "    plt.title('Target (z = ' + str(z) + ')', fontsize=15)\n",
        "\n",
        "    # Prediction\n",
        "    plt.subplot(1,4,3)\n",
        "    plt.axis('off')\n",
        "    plt.imshow(test_prediction[z-1], cmap='magma')\n",
        "    plt.title('Prediction (z = ' + str(z) + ')', fontsize=15)\n",
        "    \n",
        "    # Overlay\n",
        "    plt.subplot(1,4,4)\n",
        "    plt.axis('off')\n",
        "    plt.imshow(test_target[z-1], cmap='Greens')\n",
        "    plt.imshow(test_prediction[z-1], alpha=0.5, cmap='Purples')\n",
        "    plt.title('Overlay (z = ' + str(z) + ')', fontsize=15)\n",
        "    plt.savefig(os.path.join(qc_model_path,qc_model_name,'Quality Control')+'/QC_example_data.png', bbox_inches='tight', pad_inches=0)\n",
        "interact(scroll_in_z, z=widgets.IntSlider(min=1, max=test_source.shape[0], step=1, value=0));"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aIvRxpZlsFeZ"
      },
      "source": [
        "## **5.3. Determine best Intersection over Union and threshold**\n",
        "---\n",
        "\n",
        "<font size = 4>**Note:** This section is only relevant if the target image is a binary mask and `binary_target` is selected in Section 3!  \n",
        "\n",
        "<font size = 4>This section will provide both a visual and a quantitative indication of the model performance by comparing the overlay of the predicted and source volume, as well as computing the highest [**Intersection over Union**](https://en.wikipedia.org/wiki/Jaccard_index) (IoU) score. The IoU is also known as the Jaccard Index.  \n",
        "\n",
        "<font size = 4>The best threshold is calculated using the IoU. Each threshold value from 0 to 255 is tested and the threshold with the highest score is deemed the best. The IoU is calculated for the entire volume in 3D."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "XhkeZTFusHA8"
      },
      "outputs": [],
      "source": [
        "\n",
        "#@markdown ##Calculate Intersection over Union and best threshold \n",
        "prediction = tifffile.imread(predict_path)\n",
        "prediction = np.interp(prediction, (prediction.min(), prediction.max()), (0, 255))\n",
        "\n",
        "target = tifffile.imread(testing_target).astype(np.bool)\n",
        "\n",
        "def iou_vs_threshold(prediction, target):\n",
        "    threshold_list = []\n",
        "    IoU_scores_list = []\n",
        "\n",
        "    for threshold in range(0,256): \n",
        "        mask = prediction > threshold\n",
        "\n",
        "        intersection = np.logical_and(target, mask)\n",
        "        union = np.logical_or(target, mask)\n",
        "        iou_score = np.sum(intersection) / np.sum(union)\n",
        "\n",
        "        threshold_list.append(threshold)\n",
        "        IoU_scores_list.append(iou_score)\n",
        "\n",
        "    return threshold_list, IoU_scores_list\n",
        "\n",
        "threshold_list, IoU_scores_list = iou_vs_threshold(prediction, target)\n",
        "thresh_arr = np.array(list(zip(threshold_list, IoU_scores_list)))\n",
        "best_thresh = int(np.where(thresh_arr == np.max(thresh_arr[:,1]))[0])\n",
        "best_iou = IoU_scores_list[best_thresh]\n",
        "\n",
        "print('Highest IoU is {:.4f} with a threshold of {}'.format(best_iou, best_thresh))\n",
        "\n",
        "def adjust_threshold(threshold, z):\n",
        "\n",
        "    f=plt.figure(figsize=(25,5))\n",
        "    plt.subplot(1,4,1)\n",
        "    plt.imshow((prediction[z-1] > threshold).astype('uint8'), cmap='magma')\n",
        "    plt.title('Prediction (Threshold = ' + str(threshold) + ')', fontsize=15)\n",
        "    plt.axis('off')\n",
        "\n",
        "    plt.subplot(1,4,2)\n",
        "    plt.imshow(target[z-1], cmap='magma')\n",
        "    plt.title('Target (z = ' + str(z) + ')', fontsize=15)\n",
        "    plt.axis('off')\n",
        "\n",
        "    plt.subplot(1,4,3)\n",
        "    plt.axis('off')\n",
        "    plt.imshow(test_source[z-1], cmap='gray')\n",
        "    plt.imshow((prediction[z-1] > threshold).astype('uint8'), alpha=0.4, cmap='Reds')\n",
        "    plt.title('Overlay (z = ' + str(z) + ')', fontsize=15)\n",
        "\n",
        "    plt.subplot(1,4,4)\n",
        "    plt.title('Threshold vs. IoU', fontsize=15)\n",
        "    plt.plot(threshold_list, IoU_scores_list)\n",
        "    plt.plot(threshold, IoU_scores_list[threshold], 'ro')     \n",
        "    plt.ylabel('IoU score')\n",
        "    plt.xlabel('Threshold')\n",
        "    plt.savefig(os.path.join(qc_model_path,qc_model_name,'Quality Control')+'/QC_IoU_analysis.png',bbox_inches=matplotlib.transforms.Bbox([[17.5,0],[23,5]]),pad_inches=0)\n",
        "    plt.show()\n",
        "\n",
        "interact(adjust_threshold, \n",
        "         threshold=widgets.IntSlider(min=0, max=255, step=1, value=best_thresh),\n",
        "         z=widgets.IntSlider(min=1, max=prediction.shape[0], step=1, value=0));\n",
        "\n",
        "#Make a pdf summary of the QC results\n",
        "\n",
        "qc_pdf_export()\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S1Ks6CxfekhY"
      },
      "source": [
        "## **5.4. 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>  More information at https://deepimagej.github.io/deepimagej/"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "ZY_Y0Y4t8bL3"
      },
      "outputs": [],
      "source": [
        "# ------------- User input ------------\n",
        "# information about the model\n",
        "#@markdown ##Introduce the information to document your model:\n",
        "Trained_model_name    = \"Name of your model\" #@param {type:\"string\"}\n",
        "Trained_model_authors =  \"[Author 1 name, Author 2 name]\" #@param {type:\"string\"}\n",
        "Trained_model_authors_affiliation =  \"[Author affiliation, Author 2 affiliation]\" #@param {type:\"string\"}\n",
        "Trained_model_description = \"A 3D U-Net trained using ZeroCostDL4Mic to segment mitochondria in 3D transmission electron microscopy images of the CA1 hippocampus region.\" #@param {type:\"string\"}\n",
        "Trained_model_license = 'MIT'#@param {type:\"string\"}\n",
        "Trained_model_references = [\"Falk et al. Nature Methods 2019\", \"Ronneberger et al. arXiv in 2015\", \"Lucas von Chamier et al. biorXiv 2020\"] \n",
        "Trained_model_DOI = [\"https://doi.org/10.1038/s41592-018-0261-2\",\"https://doi.org/10.1007/978-3-319-24574-4_28\", \"https://doi.org/10.1101/2020.03.20.000133\"] \n",
        "\n",
        "# Training data\n",
        "# ---------------------------------------\n",
        "#@markdown ##Include information about training data (optional):\n",
        "include_training_data = True #@param {type: \"boolean\"}\n",
        "#@markdown ### - If it is published in the BioImage Model Zoo, please, provide the ID\n",
        "data_from_bioimage_model_zoo = False #@param {type: \"boolean\"}\n",
        "training_data_ID = ''#@param {type:\"string\"}\n",
        "#@markdown ### - If not, please provide the URL tot he data and a short description\n",
        "training_data_source = 'https://www.epfl.ch/labs/cvlab/data/data-em/'#@param {type:\"string\"}\n",
        "training_data_description = 'As said in the webpage, the dataset represents a section taken from the CA1 hippocampus region of the brain. The resolution of each voxel is approximately 5x5x5nm.'#@param {type:\"string\"}\n",
        "\n",
        "\n",
        "\n",
        "# Add example image information\n",
        "# ---------------------------------------\n",
        "#@markdown ##Choose if you want to threshold the network output and if you want to use the best threshold calculated before.\n",
        "apply_threshold = True #@param {type: \"boolean\"}\n",
        "Use_The_Best_Average_Threshold = True #@param {type:\"boolean\"}\n",
        "#@markdown ###If not, please input:\n",
        "threshold = 210 #@param {type:\"number\"}\n",
        "if Use_The_Best_Average_Threshold:\n",
        "    threshold = best_thresh\n",
        "## project it into the (0,1) range to operate with it in the bioimageio.core\n",
        "threshold = float(threshold) / 255\n",
        "\n",
        "#@markdown ##Introduce the voxel size (pixel size for each Z-slice and the distance between Z-salices) (in microns) of the image provided as an example of the model processing:\n",
        "# information about the example image\n",
        "PixelSize = 0.8 #@param {type:\"number\"}\n",
        "Zdistance = 0.005 #@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",
        "if default_example_image:\n",
        "    fileID = testing_source\n",
        "\n",
        "# Load model parameters\n",
        "# ---------------------------------------\n",
        "def last_chars(x):\n",
        "    return(x[-11:])\n",
        "try:\n",
        "    ckpt_dir_list = glob(full_qc_model_path + '/ckpt/*')\n",
        "    ckpt_dir_list.sort(key=last_chars)\n",
        "    last_ckpt_path = ckpt_dir_list[0]\n",
        "    print('Predicting from checkpoint:', os.path.basename(last_ckpt_path))\n",
        "except IndexError:\n",
        "    raise CheckpointError('No previous checkpoints were found, please retrain model.')\n",
        "\n",
        "params = pd.read_csv(os.path.join(full_qc_model_path, 'params.csv'), names=['val'], header=0, index_col=0)   \n",
        "# Load the model and process the example image\n",
        "# ---------------------------------------\n",
        "model = Unet3D(shape=params.loc['training_shape', 'val'])\n",
        "model.model.load_weights(last_ckpt_path)\n",
        "# ------------- Execute bioimage model zoo configuration ------------\n",
        "# Create a model without compilation so it can be used in any other environment.\n",
        "unet = model.model\n",
        "input  = unet.input\n",
        "single_output = unet.output\n",
        "# remove the custom loss function from the model, so that it can be used outside of this notebook\n",
        "unet = Model(input, single_output)\n",
        "weight_path = os.path.join(full_qc_model_path, 'keras_weights.hdf5')\n",
        "unet.save(weight_path)\n",
        "\n",
        "# create the author spec input\n",
        "auth_names = Trained_model_authors[1:-1].split(\",\")\n",
        "auth_affs = Trained_model_authors_affiliation[1:-1].split(\",\")\n",
        "assert len(auth_names) == len(auth_affs)\n",
        "authors = [{\"name\": auth_name, \"affiliation\": auth_aff} for auth_name, auth_aff in zip(auth_names, auth_affs)]\n",
        "\n",
        "# I would recommend using CCBY-4 as licencese\n",
        "license = Trained_model_license\n",
        "\n",
        "# where to save the model\n",
        "output_root = os.path.join(full_qc_model_path, Trained_model_name + '.bioimage.io.model')\n",
        "os.makedirs(output_root, exist_ok=True)\n",
        "output_path = os.path.join(output_root, f\"{Trained_model_name}.zip\")\n",
        "\n",
        "# create a markdown readme with information\n",
        "readme_path = os.path.join(output_root, \"README.md\")\n",
        "with open(readme_path, \"w\") as f:\n",
        "  f.write(\"Visit https://github.com/HenriquesLab/ZeroCostDL4Mic/wiki\")\n",
        "\n",
        "# create the citation input spec\n",
        "assert len(Trained_model_DOI) == len(Trained_model_references)\n",
        "citations = [{'text': text, 'doi': doi} for text, doi in zip(Trained_model_references, Trained_model_DOI)]\n",
        "\n",
        "# create the training data\n",
        "if include_training_data:\n",
        "    if data_from_bioimage_model_zoo:\n",
        "      training_data = {\"id\": training_data_ID}\n",
        "    else:\n",
        "      training_data = {\"source\": training_data_source,\n",
        "                       \"description\": training_data_description}\n",
        "else:\n",
        "    training_data={}\n",
        "\n",
        "# create the input spec\n",
        "min_percentile = 0\n",
        "max_percentile = 100\n",
        "shape = [sh for sh in unet.input.shape]\n",
        "# batch should never be constrained\n",
        "assert shape[0] is None\n",
        "shape[0] = 1  # batch is set to 1 for bioimage.io\n",
        "assert all(sh is not None for sh in shape)  # make sure all other shapes are fixed\n",
        "pixel_size = {\"x\": PixelSize, \"y\": PixelSize, 'z': Zdistance}\n",
        "kwargs = dict(\n",
        "  input_names=[\"input\"],\n",
        "  input_axes=[\"bxyzc\"],\n",
        "  pixel_sizes=[pixel_size],\n",
        "  preprocessing=[[{\"name\": \"scale_range\", \"kwargs\": {\"min_percentile\": min_percentile, \n",
        "                                  \"max_percentile\": max_percentile, \n",
        "                                  \"mode\": \"per_sample\",\n",
        "                                  \"axes\": \"xyzc\"}}]])\n",
        "\n",
        "shape = tuple(shape)\n",
        "\n",
        "if apply_threshold:\n",
        "  print(\"The model output is thresholded\")\n",
        "  postprocessing = [[{\"name\": \"scale_range\", \"kwargs\": {\"min_percentile\": 0, \n",
        "                                  \"max_percentile\": 100, \n",
        "                                  \"mode\": \"per_sample\",\n",
        "                                  \"axes\": \"xyzc\"}}, \n",
        "                     {\"name\": \"binarize\", \"kwargs\": {\"threshold\": threshold}}]]\n",
        "else:\n",
        "  print(\"The model output is not thresholded\")\n",
        "  postprocessing = None\n",
        "\n",
        "output_spec = dict(\n",
        "  output_names=[\"output\"],\n",
        "  output_axes=[\"bxyzc\"],\n",
        "  postprocessing=postprocessing,\n",
        "  output_reference=[\"input\"],\n",
        "  output_scale=[5*[1]], # consider changing it if the input has more than one channel\n",
        "  output_offset=[5*[0]]\n",
        ")\n",
        "kwargs.update(output_spec)\n",
        "\n",
        "# load the input image, crop it if necessary, and save as numpy file\n",
        "# The crop will be centered to get an image with some content.\n",
        "test_img = tifffile.imread(fileID).T  \n",
        "\n",
        "\n",
        "x_size = int(test_img.shape[0]/2)\n",
        "x_size = x_size-int(shape[1]/2)\n",
        "\n",
        "y_size = int(test_img.shape[1]/2)\n",
        "y_size = y_size-int(shape[2]/2)\n",
        "\n",
        "z_size = int(test_img.shape[2]/2)\n",
        "z_size = z_size-int(shape[3]/2)\n",
        "\n",
        "assert test_img.ndim == 3\n",
        "test_img = test_img[x_size : x_size + shape[1],\n",
        "                    y_size : y_size + shape[2],\n",
        "                    z_size : z_size + shape[3]]\n",
        "\n",
        "assert test_img.shape == shape[1:4], f\"{test_img.shape}, {shape}\"\n",
        "# Save the test image\n",
        "test_in_path = os.path.join(output_root, \"test_input.npy\")\n",
        "test_img = test_img.astype(np.float32)\n",
        "np.save(test_in_path, test_img[None, ..., None])  # add batch and channel axis\n",
        "# Normalize the image before adding batch and channel dimensions\n",
        "eps = 1e-6 # Same epsilon value as in the bioimageio.core library\n",
        "n = test_img - np.min(test_img)\n",
        "d = np.max(test_img) - np.min(test_img) + eps\n",
        "test_img = n/d\n",
        "test_img = test_img[None, ..., None]\n",
        "test_prediction = unet.predict(test_img)\n",
        "\n",
        "# run prediction on the input image and save the result as expected output\n",
        "if apply_threshold:\n",
        "  # Normalise the values of the test image\n",
        "\n",
        "  n = test_prediction - np.min(test_prediction)\n",
        "  d = np.max(test_prediction) - np.min(test_prediction) + 1e-6\n",
        "  test_prediction = n/d\n",
        "\n",
        "  test_prediction = (np.squeeze(test_prediction) > threshold).astype(np.uint8)\n",
        "else:\n",
        "  test_prediction = np.squeeze(test_prediction)\n",
        "assert test_prediction.ndim == 3\n",
        "test_prediction = test_prediction[None, ..., None]  # add batch and channel axis\n",
        "test_out_path = os.path.join(output_root, \"test_output.npy\")\n",
        "np.save(test_out_path, test_prediction)\n",
        "\n",
        "# attach the QC report to the model (if it exists)\n",
        "qc_path = os.path.join(full_qc_model_path, 'Quality Control', 'training_evaluation.csv')\n",
        "if os.path.exists(qc_path):\n",
        "  attachments = {\"files\": [qc_path]}\n",
        "else:\n",
        "  attachments = None\n",
        "\n",
        "# export the model with keras weihgts\n",
        "build_model(\n",
        "    weight_uri=weight_path,\n",
        "    test_inputs=[test_in_path],\n",
        "    test_outputs=[test_out_path],\n",
        "    name=Trained_model_name,\n",
        "    description=Trained_model_description,\n",
        "    authors=authors,\n",
        "    tags=['zerocostdl4mic', 'deepimagej', 'segmentation', 'tem', 'unet'],\n",
        "    license=license,\n",
        "    documentation=readme_path,\n",
        "    cite=citations,\n",
        "    output_path=output_path,\n",
        "    add_deepimagej_config=True,\n",
        "    tensorflow_version=tf.__version__,\n",
        "    attachments=attachments,\n",
        "    training_data = training_data,\n",
        "    **kwargs\n",
        ")\n",
        "\n",
        "# convert the keras weights to tensorflow and add them to the model\n",
        "tf_weight_path = os.path.join(full_qc_model_path, \"tf_weights\")\n",
        "# we need to make sure that the tf weight folder does not exist\n",
        "if os.path.exists(tf_weight_path):\n",
        "  rmtree(tf_weight_path)\n",
        "convert_weights_to_tensorflow_saved_model_bundle(output_path, tf_weight_path + \".zip\")\n",
        "add_weights(output_path, tf_weight_path + \".zip\", output_path, tensorflow_version=tf.__version__)\n",
        "\n",
        "# check that the model works for keras and tensorflow \n",
        "res = test_model(output_path, weight_format=\"keras_hdf5\")\n",
        "success = True\n",
        "if res[-1][\"error\"] is not None:\n",
        "  success = False\n",
        "  print(\"test-model failed for keras weights:\", res[-1][\"error\"])\n",
        "  \n",
        "res = test_model(output_path, weight_format=\"tensorflow_saved_model_bundle\")\n",
        "if res[-1][\"error\"] is not None:\n",
        "  success = False\n",
        "  print(\"test-model failed for tensorflow weights:\", res[-1][\"error\"])\n",
        "\n",
        "if success:\n",
        "  print(\"The bioimage.io model was successfully exported to\", output_path)\n",
        "else:\n",
        "  print(\"The bioimage.io model was exported to\", output_path)\n",
        "  print(\"Some tests of the model did not work! You can still download and test the model.\")\n",
        "  print(\"You can still download and test the model, but it may not work as expected.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-tJeeJjLnRkP"
      },
      "source": [
        "# **6. Using the trained model**\n",
        "\n",
        "---\n",
        "\n",
        "<font size = 4>Once sufficient performance of the trained model has been established using Section 5, the network can be used to segment unseen volumetric data."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d8wuQGjoq6eN"
      },
      "source": [
        "## **6.1. Generate predictions from unseen dataset**\n",
        "---\n",
        "\n",
        "<font size = 4>The most recently trained model can now be used to predict segmentation masks on unseen images. If you want to use an older model, leave `model_path`  blank. Predicted output images are saved in `output_path` as Image-J compatible TIFF files.\n",
        "\n",
        "## **Prediction parameters**\n",
        "\n",
        "* <font size = 4>**`source_path`** specifies the location of the source \n",
        "image volume.\n",
        "\n",
        "* <font size = 4>**`output_directory`** specified the directory where the output predictions are stored.\n",
        "\n",
        "* <font size = 4>**`binary_target`** should be chosen if the network is trained to predict binary segmentation masks.\n",
        "\n",
        "* <font size = 4>**`threshold`** can be calculated in Section 5 and is used to generate binary masks from the predictions.\n",
        "\n",
        "* <font size = 4>**`big_tiff`** should be chosen if the expected prediction exceeds 4GB. The predictions will be saved using the BigTIFF format. Beware that this might substantially reduce the prediction speed. *Default: False* \n",
        "\n",
        "* <font size = 4>**`prediction_depth`** is only relevant if the prediction is saved as a BigTIFF. The prediction will not be performed in one go to not deplete the memory resources. Instead, the prediction is iteratively performed on a subset of the entire volume  with shape `(source.shape[0], source.shape[1], prediction_depth)`. *Default: 32*\n",
        "\n",
        "* <font size = 4>**`model_path`** specifies the path to a model other than the most  recently trained."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "DEmhPh5fsWX2"
      },
      "outputs": [],
      "source": [
        "#@markdown ## Download example volume\n",
        "\n",
        "#@markdown <font size = 4> This can take up to an hour\n",
        "\n",
        "import requests  \n",
        "import os\n",
        "from tqdm.notebook import tqdm \n",
        "\n",
        "\n",
        "def download_from_url(url, save_as):\n",
        "    file_url = url\n",
        "    r = requests.get(file_url, stream=True)  \n",
        "  \n",
        "    with open(save_as, 'wb') as file:  \n",
        "        for block in tqdm(r.iter_content(chunk_size = 1024), desc = 'Downloading ' + os.path.basename(save_as), total=3275073, ncols=1000):\n",
        "            if block:\n",
        "                file.write(block)  \n",
        "\n",
        "download_from_url('https://documents.epfl.ch/groups/c/cv/cvlab-unit/www/data/%20ElectronMicroscopy_Hippocampus/volumedata.tif', 'example_dataset/volumedata.tif')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "y2TD5p7MZrEb"
      },
      "outputs": [],
      "source": [
        "#@markdown ### Provide the path to your dataset and to the folder where the predictions are saved, then run the cell to predict outputs from your unseen images.\n",
        "\n",
        "source_path = \"\" #@param {type:\"string\"}\n",
        "output_directory = \"\" #@param {type:\"string\"}\n",
        "\n",
        "if not os.path.exists(output_directory):\n",
        "    os.makedirs(output_directory)\n",
        "\n",
        "output_path = os.path.join(output_directory, os.path.splitext(os.path.basename(source_path))[0] + '_predicted.tif')\n",
        "#@markdown ###Prediction parameters:\n",
        "\n",
        "binary_target = True #@param {type:\"boolean\"}\n",
        "\n",
        "save_probability_map = False #@param {type:\"boolean\"}\n",
        "\n",
        "#@markdown <font size = 3>Determine best threshold in Section 5.2.\n",
        "\n",
        "use_calculated_threshold = True #@param {type:\"boolean\"}\n",
        "threshold =  200#@param {type:\"number\"}\n",
        "\n",
        "# Tifffile library issues means that images cannot be appended to \n",
        "#@markdown <font size = 3>Choose if prediction file exceeds 4GB or if input file is very large (above 2GB). Image volume saved as BigTIFF.\n",
        "big_tiff = False #@param {type:\"boolean\"}\n",
        "\n",
        "#@markdown <font size = 3>Reduce `prediction_depth` if runtime runs out of memory during prediction. Only relevant if prediction saved as BigTIFF\n",
        "\n",
        "prediction_depth =  32#@param {type:\"number\"}\n",
        "\n",
        "#@markdown ###Model to be evaluated\n",
        "#@markdown <font size = 3>If left blank, the latest model defined in Section 5 will be evaluated\n",
        "\n",
        "full_model_path_ = \"\" #@param {type:\"string\"}\n",
        "\n",
        "if len(full_model_path_) == 0:\n",
        "    full_model_path_ = os.path.join(qc_model_path, qc_model_name) \n",
        "\n",
        "\n",
        "\n",
        "# Load parameters\n",
        "params = pd.read_csv(os.path.join(full_model_path_, 'params.csv'), names=['val'], header=0, index_col=0)   \n",
        "model = Unet3D(shape=params.loc['training_shape', 'val'])\n",
        "\n",
        "if use_calculated_threshold:\n",
        "    threshold = best_thresh\n",
        "\n",
        "def last_chars(x):\n",
        "    return(x[-11:])\n",
        "\n",
        "try:\n",
        "    ckpt_dir_list = glob(full_model_path_ + '/ckpt/*')\n",
        "    ckpt_dir_list.sort(key=last_chars)\n",
        "    last_ckpt_path = ckpt_dir_list[0]\n",
        "    print('Predicting from checkpoint:', os.path.basename(last_ckpt_path))\n",
        "except IndexError:\n",
        "    raise CheckpointError('No previous checkpoints were found, please retrain model.')\n",
        "\n",
        "src = tifffile.imread(source_path)\n",
        "\n",
        "if src.nbytes >= 4e9:\n",
        "    big_tiff = True\n",
        "    print('The source file exceeds 4GB in memory, prediction will be saved as BigTIFF!')\n",
        "\n",
        "if binary_target:\n",
        "    if not big_tiff:\n",
        "        prediction = model.predict(src, last_ckpt_path, downscaling=params.loc['downscaling', 'val'], true_patch_size=params.loc['true_patch_size', 'val'])\n",
        "        prediction = np.interp(prediction, (prediction.min(), prediction.max()), (0, 255))\n",
        "        prediction = (prediction > threshold).astype('float32')\n",
        "\n",
        "        tifffile.imwrite(output_path, prediction, imagej=True)\n",
        "\n",
        "    else:\n",
        "        with tifffile.TiffWriter(output_path, bigtiff=True) as tif:\n",
        "            for i in tqdm(range(0, src.shape[0], prediction_depth)):\n",
        "                prediction = model.predict(src, last_ckpt_path, z_range=(i,i+prediction_depth), downscaling=params.loc['downscaling', 'val'], true_patch_size=params.loc['true_patch_size', 'val'])\n",
        "                prediction = np.interp(prediction, (prediction.min(), prediction.max()), (0, 255))\n",
        "                prediction = (prediction > threshold).astype('float32')\n",
        "                \n",
        "                for j in range(prediction.shape[0]):\n",
        "                    tif.save(prediction[j])\n",
        "\n",
        "if not binary_target or save_probability_map:\n",
        "    if not binary_target:\n",
        "        prob_map_path = output_path\n",
        "    else:\n",
        "        prob_map_path = os.path.splitext(output_path)[0] + '_prob_map.tif'\n",
        "    \n",
        "    if not big_tiff:\n",
        "        prediction = model.predict(src, last_ckpt_path, downscaling=params.loc['downscaling', 'val'], true_patch_size=params.loc['true_patch_size', 'val'])\n",
        "        prediction = np.interp(prediction, (prediction.min(), prediction.max()), (0, 255))\n",
        "        tifffile.imwrite(prob_map_path, prediction.astype('float32'), imagej=True)\n",
        "\n",
        "    else:\n",
        "        with tifffile.TiffWriter(prob_map_path, bigtiff=True) as tif:\n",
        "            for i in tqdm(range(0, src.shape[0], prediction_depth)):\n",
        "                prediction = model.predict(src, last_ckpt_path, z_range=(i,i+prediction_depth), downscaling=params.loc['downscaling', 'val'], true_patch_size=params.loc['true_patch_size', 'val'])\n",
        "                prediction = np.interp(prediction, (prediction.min(), prediction.max()), (0, 255))\n",
        "                \n",
        "                for j in range(prediction.shape[0]):\n",
        "                    tif.save(prediction[j])\n",
        "\n",
        "print('Predictions saved as', output_path)\n",
        "\n",
        "src_volume = tifffile.imread(source_path)\n",
        "pred_volume = tifffile.imread(output_path)\n",
        "\n",
        "def scroll_in_z(z):\n",
        "  \n",
        "    f=plt.figure(figsize=(25,5))\n",
        "    plt.subplot(1,2,1)\n",
        "    plt.imshow(src_volume[z-1], cmap='gray')\n",
        "    plt.title('Source (z = ' + str(z) + ')', fontsize=15)\n",
        "    plt.axis('off')\n",
        "\n",
        "    plt.subplot(1,2,2)\n",
        "    plt.imshow(pred_volume[z-1], cmap='magma')\n",
        "    plt.title('Prediction (z = ' + str(z) + ')', fontsize=15)\n",
        "    plt.axis('off')\n",
        "\n",
        "interact(scroll_in_z, z=widgets.IntSlider(min=1, max=src_volume.shape[0], step=1, value=0));\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hvkd66PldsXB"
      },
      "source": [
        "## **6.2. Download your predictions**\n",
        "---\n",
        "\n",
        "<font size = 4>**Store your data** and ALL its results elsewhere by downloading it from Google Drive and after that clean the original folder tree (datasets, results, trained model etc.) if you plan to train or use new networks. Please note that the notebook will otherwise **OVERWRITE** all files which have the same name."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "q3lSeWp3G8eD"
      },
      "source": [
        "# **7. Version log**\n",
        "\n",
        "---\n",
        "<font size = 4>**v2.1**: \n",
        "* Updated notebook to use with TensorFlow 2.11\n",
        "* The full requirements file is exported to replicate the environment\n",
        "* Bug in the preprocessing to be compatible with the BioImage Model Zoo is solved.\n",
        "* Installs `fpdf2`, curates the format of the pdf and fixes bugs in `keras` versioning writting.\n",
        "\n",
        "---"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UvSlTaH14s3t"
      },
      "source": [
        "\n",
        "#**Thank you for using 3D U-Net!**"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "machine_shape": "hm",
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "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.11.0"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}