DAIN_4_Microscopy.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "DAIN_4_Microscopy.ipynb",
      "provenance": [],
      "collapsed_sections": [
        "C-wdtVN5KUFi",
        "S0lI6iAOLVLd",
        "2B7v00leuXQM",
        "yRGaHe8KAzSJ",
        "0XUazMIVI_8w",
        "AMkubWbUJCsr",
        "Z7SbuQBjtG93",
        "9jlVnNSVUM4p",
        "UKeq9iSVz6sM",
        "IWFtMREDS_LB",
        "WT2RF0xWVWJm",
        "d_Ov3NRSJ-gw",
        "MX4N4Md6bHZF",
        "wEmV2hTZbT-z",
        "2Hq0tP5bXWCS",
        "tta4phe6VMtO",
        "4_eedldMK8bs"
      ]
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "C-wdtVN5KUFi"
      },
      "source": [
        "#**DAIN 4 Microscopy**\n",
        "[**TUTORIAL Video **](https://youtu.be/RyMQuRYtpbM): **WATCH THAT BEFORE USING IT**!!!\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S0lI6iAOLVLd"
      },
      "source": [
        "##Content Aware Frame Interpolation (CAFI) in T or/and Z dimension for Microscopy Images \n",
        "\n",
        "---\n",
        "\n",
        "\n",
        "<font size = 4> \"DAIN\" is a neural network architecture capable of interpolating images along the T and/or Z-dimension of microscopy image series. This particular notebook allows the user to perform a frame interpolation of **3D+t or 2D+t** microscopy images in **grayscale or RGB** which means it increases the **frame rate in T-and Z-dimension by a chosen factor between 2 and 4** \n",
        "\n",
        "---\n",
        "<font size = 4> Example for illustration: Image input dimensions with factor 2 improvements: (t,z,x,y) 10:15:256:256 -> Image output dimensions: 19:29:256:256\n",
        "Image input dimensions with factor 3 improvements: (t,z,x,y) 10:15:256:256 -> Image output dimensions: 28:43:256:256\n",
        "\n",
        "---\n",
        "\n",
        "<font size = 4>This can be achieved by using a pretrained model or/and by finetuning our pretrained model with your own dataset.\n",
        "\n",
        "\n",
        "\n",
        "The network used in this application is using a depth-aware flow projection with flow estimation, depth estimation and context extraction to create the interpolated images and achieves state of the art quantitative and qualitative performance in the image interpolation task. \n",
        "\n",
        "---\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "w1OMrPQPqY6y"
      },
      "source": [
        "---\n",
        "<font size = 4> *Disclaimer*:\n",
        "\n",
        "<font size = 4> This notebook is provided as an addon to the the Zero-Cost Deep-Learning to Enhance Microscopy project (https://github.com/HenriquesLab/DeepLearning_Collab/wiki). Jointly developed by the Jacquemet (link to https://cellmig.org/) and Henriques (https://henriqueslab.github.io/) laboratories.\n",
        "\n",
        "<font size = 4>  The network was first published in the context of space-time video super-resolution in 2019 under the name: **Depth-Aware Video Frame Interpolation** by *Wenbo Bao, Wei-Sheng Lai, Chao Ma, Xiaoyun Zhang, Zhiyong Gao, and Ming-Hsuan Yang* [URL](https://arxiv.org/abs/1904.00830)\n",
        "\n",
        "[URL](https://openaccess.thecvf.com/content_CVPR_2019/papers/Bao_Depth-Aware_Video_Frame_Interpolation_CVPR_2019_paper.pdf) IEEE Conference on Computer Vision and Pattern Recognition (CVPR 2019).\n",
        "\n",
        "<font size = 4> The original source code used for video-frame interpolation can be found in:\n",
        "\n",
        "[URL](https://github.com/baowenbo/DAIN) (in this Colab notebook we use a modified version of this repository which can be found here [URL](https://github.com/mpriessner/VFIN)\n",
        "\n",
        "<font size = 4> **Please also cite this original paper when using or developing this notebook.** \n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2B7v00leuXQM"
      },
      "source": [
        "## **How to use this notebook?**\n",
        "---\n",
        "\n",
        "<font size = 4>Video describing how to use our notebooks are available on youtube:\n",
        "  - [**Video 1**](https://www.youtube.com/watch?v=GzD2gamVNHI&feature=youtu.be): Full run through of the workflow to obtain the notebooks and the provided test datasets as well as a common use of the notebook\n",
        "  - [**Video 2**](https://www.youtube.com/watch?v=PUuQfP5SsqM&feature=youtu.be): Detailed description of the different sections of the notebook\n",
        "  \n",
        "  - [**Video 3**](https://www.youtube.com/watch?v=PUuQfP5SsqM&feature=youtu.be): This video goes through the complete workflow of this specific workbook **WATCH THAT BEFORE USING IT**\n",
        "\n",
        "\n",
        "---"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_tvQ2uO5ugVv"
      },
      "source": [
        "###**Structure of a notebook**\n",
        "\n",
        "<font size = 4>The notebook contains two types of cell:  \n",
        "\n",
        "<font size = 4>**Text cells** provide information and can be modified by douple-clicking the cell. You are currently reading the text cell. You can create a new text by clicking `+ Text`.\n",
        "\n",
        "<font size = 4>**Code cells** contain code and the code can be modfied by selecting the cell. To execute the cell, move your cursor on the `[ ]`-mark on the left side of the cell (play button appears). Click to execute the cell. After execution is done the animation of play button stops. You can create a new coding cell by clicking `+ Code`.\n",
        "\n",
        "---\n",
        "###**Table of contents, Code snippets** and **Files**\n",
        "\n",
        "<font size = 4>On the top left side of the notebook you find three tabs which contain from top to bottom:\n",
        "\n",
        "<font size = 4>*Table of contents* = contains structure of the notebook. Click the content to move quickly between sections.\n",
        "\n",
        "<font size = 4>*Code snippets* = contain examples how to code certain tasks. You can ignore this when using this notebook.\n",
        "\n",
        "<font size = 4>*Files* = contain all available files. After mounting your google drive (see section 0) you will find your files and folders here. \n",
        "\n",
        "<font size = 4>**Remember that all uploaded files are purged after changing the runtime.** All files saved in Google Drive will remain. You do not need to use the Mount Drive-button; your Google Drive is connected in section 0.\n",
        "\n",
        "If you demand bigger amounts of data or quicker loading speed you can also connect to your google cloud bucket by providing your `project_id` and `bucket_name`\n",
        "\n",
        "The advantage of google cloud buckets is that they don't have any data volume restrictions that might become a problem for big datasets that need to be accessed via google drive. The service  for data-storage is fairly cheap and might be worth considering in case one runs into this specific problem. \n",
        "\n",
        "<font size = 4>**Note:** The \"sample data\" in \"Files\" contains default files. Do not upload anything in here!\n",
        "\n",
        "---\n",
        "###**Making changes to the notebook**\n",
        "\n",
        "<font size = 4>**You can make a copy** of the notebook and save it to your Google Drive. To do this click file -> save a copy in drive.\n",
        "\n",
        "<font size = 4>To **edit a cell**, double click on the text. This will show you either the source code (in code cells) or the source text (in text cells).\n",
        "You can use the `#`-mark in code cells to comment out parts of the code. This allows you to keep the original code piece in the cell as a comment."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yRGaHe8KAzSJ"
      },
      "source": [
        "#**0.0 Before getting started**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1gVXhZlqw1Oi"
      },
      "source": [
        "\n",
        "---\n",
        "<font size = 4> This notebook provides the opportunity to either use the network with the demonstration dataset provided from the paper or to use your own training data to test the algorithm on your own data.\n",
        "\n",
        "<font size = 4> The notebook may require a large amount of disk space. If training your own network with your own datasets, the available disk space on the user's google drive should contain at least 5-10GB.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f5_i9KDwzpRP"
      },
      "source": [
        "---\n",
        "<font size = 4>**Data Format**\n",
        "\n",
        "**The data used to train or test the DAIN network can be the following: 3D (TYX), 4D (TZYX or TYXC) or 5D image stacks (TZYXC) in .tiff (.tif) file format.**\n",
        "\n",
        "To use this notebook on user data, upload the data in the your google drive and execute the notebook cells step by step following the instructions in the associated section. \n",
        "\n",
        "<font size = 4>- If you wish to **Train or finetune your model** or **perform frame or slice interpolation** using a model previously generated, you will first need to run **Sections 2** to set up the notebook.\n",
        "\n",
        "<font size = 4>- If you wish to **run predictions/interpolation** using a model previously generated and saved on your Google Drive, you will  need to run **Sections 1 and 2** to set up the notebook, then use **Section 5** to run the predictions.\n",
        "\n",
        "<font size = 4>- If you plan to train or finetune the network on a very big dataset (more than 10 GB) we recommend to perform the training data preparation with offline on your PC running the python scrip provided in the following folder: `/content/DAIN/script` \n",
        "\n",
        "<font size = 4>- If you want to **check the quality** of the network first run the **downsampling 3** and perform the image interpolation on the downsampled. The result can then be compared with the ground truth image in **Section 6**. \n",
        "\n",
        "\n",
        "<font size = 4>**Important note**\n",
        "\n",
        "- If when running the first cell the following error occurs \"not authored by google\" just press \"Run anyways\"\n",
        "\n",
        "- If for some reason the notebook crashes and you cannot continue at all, you can first try to restart the runtime at `Runtime` -> `Restart Runtime`. This will delete all the variables and will reload all the libraries. (No re-install required). If that doesn't work you can perform a factory reset   `Runtime` -> `Factory reset runtime`. This will start a new virtual machine in the cloud and requires a new installation of the tool. \n",
        "\n",
        "- Make sure the training dataset is not more than around 10 GB otherwise google Colab might throw Error [5]Input/output error which means that you used up your quota (data-tranfer-volume for Google Drive). To overcome this problem we also provided a connection option to Google Cloud Bucket which can be used without limits and low costs.\n",
        "\n",
        "- Training/Finetuning needs to be completed within 12h (for free Colab) 24h (for Colab Pro) - therefore it is recommended to do smaller training batches and keep continuing training the network. A copy of the pretrained network folder will be saved in a selected folder on your Google Drive. \n",
        "\n",
        "- Your *dataset_folder* should not have spaces or brackets in its name as this is not recognized by the code and will throw an error \n",
        "\n",
        "- When first connecting the notebook there will be warning saying \"not authored by google\" which can be ignored and just click \"connect\" to start with the notebook \n",
        "\n",
        "---\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0XUazMIVI_8w"
      },
      "source": [
        "# **1. Initialise the Colab session**\n",
        "---"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AMkubWbUJCsr"
      },
      "source": [
        "\n",
        "### **1.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>**Accelator: GPU** *(Graphics processing unit)*\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "b5cpe695JCL0",
        "cellView": "form"
      },
      "source": [
        "#@markdown ##Run this cell to check if you have GPU access\n",
        "%tensorflow_version 1.x\n",
        "\n",
        "import tensorflow as tf\n",
        "if tf.test.gpu_device_name()=='':\n",
        "  print('You do not have GPU access.') \n",
        "  print('Did you change your runtime ?') \n",
        "  print('If the runtime settings are correct then Google did not allocate GPU to your session')\n",
        "  print('Expect slow performance. To access GPU try reconnecting later')\n",
        "\n",
        "else:\n",
        "  print('You have GPU access')\n",
        "\n",
        "from tensorflow.python.client import device_lib \n",
        "device_lib.list_local_devices()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Z7SbuQBjtG93"
      },
      "source": [
        "### **1.2. Mount your Google Drive or Google Cloud Bucket**\n",
        "---\n",
        "<font size = 4> To use this notebook on the data present in your Google Drive,  or your Google Cloud Bucket you need to mount your Google Drive or Google Cloud Bucket to this notebook.\n",
        "\n",
        "<font size = 4> Play one of the two the cells below to mount your Google Drive  or Google Cloud Bucket and follow the link. In the new browser window, select your drive and select 'Allow', copy the code, paste into the cell and press enter. This will give Colab access to the data on the drive. For Google Cloud Bucket you need to provide your `bucket_name` and `project_id`.\n",
        "\n",
        "<font size = 4> Once this is done, your data are available in the **Files** tab on the top left of notebook."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BZmk2w6dJSYm",
        "cellView": "form",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "0932e981-7aad-4194-eb50-7f8570b750a7"
      },
      "source": [
        "#@markdown ##Play the cell to connect your Google Drive to Colab\n",
        "\n",
        "#@markdown * Click on the URL. \n",
        "\n",
        "#@markdown * Sign in your Google Account. \n",
        "\n",
        "#@markdown * Copy the authorization code. \n",
        "\n",
        "#@markdown * Enter the authorization code. \n",
        "\n",
        "#@markdown * Click on \"Files\" site on the right. Refresh the site. Your Google Drive folder should now be available here as \"drive\". \n",
        "\n",
        "# mount user's Google Drive to Google Colab.\n",
        "from google.colab import drive\n",
        "drive.mount('/content/gdrive')\n",
        "from IPython.display import clear_output\n"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Mounted at /content/gdrive\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9jlVnNSVUM4p"
      },
      "source": [
        "### Alternative for Google Drive - Google Cloud Storage (paid accound necessary) \n",
        "Expand to connect to Google Cloud Bucket "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5nSrYOgAPwqu",
        "cellView": "form"
      },
      "source": [
        "#@markdown ##Play the cell to connect your Google Cloud Bucket\n",
        "\n",
        "#@markdown * Click on the URL. \n",
        "\n",
        "#@markdown * Sign in your Google Account. \n",
        "\n",
        "#@markdown * Copy the authorization code. \n",
        "\n",
        "#@markdown * Enter the authorization code. \n",
        "\n",
        "#@markdown * Click on \"Files\" site on the right. Refresh the site. Your Google Drive folder should now be available here as \"drive\". \n",
        "\n",
        "# mount user's Cloud to Google Colab.\n",
        "import uuid\n",
        "from google.colab import auth\n",
        "from IPython.display import clear_output\n",
        "import os\n",
        "os.chdir(\"/content\")\n",
        "\n",
        "auth.authenticate_user()\n",
        "project_id = \"\" #@param {type:\"string\"}\n",
        "bucket_name = \"\" #@param {type:\"string\"}\n",
        "bucket_name = bucket_name + str(uuid.uuid1())\n",
        "!gcloud config set project {project_id}\n",
        "!echo \"deb http://packages.cloud.google.com/apt gcsfuse-bionic main\" > /etc/apt/sources.list.d/gcsfuse.list\n",
        "!curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key add -\n",
        "!apt -qq update\n",
        "!apt -qq install gcsfuse\n",
        "!mkdir Google_Cloud\n",
        "!gcsfuse --implicit-dirs martin_phd_project_data /content/Google_Cloud\n",
        "# clear_output()\n",
        "#@markdown Before running the code make sure that you change the name of your bucket in the second last line of this cell: `!gcsfuse --implicit-dirs your_bucket_name /content/Google_Cloud`\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UKeq9iSVz6sM"
      },
      "source": [
        "\n",
        "# **2. Install DAIN network and dependencies**\n",
        "---\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VE1Xv7zhuodB",
        "cellView": "form"
      },
      "source": [
        "#@title ### Install the requirements \n",
        "#@markdown By running the cells the system requirements, libraries and the network are downloaded and installed .\n",
        "#@markdown This whole process may take 5-10 minutes and **will automatically restart the runtime** after finishing. \n",
        "import os\n",
        "\n",
        "!wget -c https://repo.anaconda.com/miniconda/Miniconda3-4.5.4-Linux-x86_64.sh\n",
        "!chmod +x Miniconda3-4.5.4-Linux-x86_64.sh\n",
        "!bash ./Miniconda3-4.5.4-Linux-x86_64.sh -b -f -p /usr/local\n",
        "!conda install pytorch==1.1 cudatoolkit torchvision -c pytorch -y  # does it with cuda 10 TRAINING AND TESTING works\n",
        "!conda install ipykernel -y\n",
        "!conda install -c conda-forge opencv -y\n",
        "\n",
        "!yes y | pip uninstall scipy\n",
        "!pip install scipy==1.1.0\n",
        "!pip install imageio\n",
        "\n",
        "### use this instead\n",
        "# https://allencellmodeling.github.io/aicsimageio/aicsimageio.html\n",
        "# https://github.com/AllenCellModeling/aicsimageio\n",
        "!pip install aicsimageio==3.2.3\n",
        "!pip install aicsimageprocessing\n",
        "!pip install tifffile \n",
        "!pip install pympler\n",
        "\n",
        "### Download Git repo\n",
        "import shutil\n",
        "# !git clone https://github.com/mpriessner/VFIN.git DAIN2\n",
        "!git clone https://github.com/mpriessner/CAFI.git\n",
        "shutil.move(\"/content/CAFI/DAIN4Mic\", \"/content/DAIN\")\n",
        "shutil.rmtree(\"/content/CAFI\")\n",
        "\n",
        "#Freate folder and download pretained models\n",
        "import os\n",
        "if not os.path.exists(\"/content/DAIN/MegaDepth/checkpoints/test_local\"):\n",
        "  os.mkdir(\"/content/DAIN/MegaDepth/checkpoints\")\n",
        "  os.mkdir(\"/content/DAIN/MegaDepth/checkpoints/test_local\")\n",
        "\n",
        "os.chdir(\"/content/DAIN/MegaDepth/checkpoints/test_local\")\n",
        "!wget http://vllab1.ucmerced.edu/~wenbobao/DAIN/best_generalization_net_G.pth\n",
        "os.chdir(\"/content/DAIN/PWCNet\")\n",
        "!wget http://vllab1.ucmerced.edu/~wenbobao/DAIN/pwc_net.pth.tar\n",
        "os.chdir(\"/content/DAIN\")\n",
        "\n",
        "### Install DAIN \n",
        "%cd /content/DAIN\n",
        "!python build.py  -cc 37,61,60,70,75  # this was the reason why it was not working before\n",
        "%cd ..\n",
        "\n",
        "### Download pretrained model\n",
        "import sys\n",
        "import os\n",
        "_ = (sys.path.append(\"/usr/local/lib/python3.6/site-packages\"))\n",
        "sys.path.insert(0,'/content/DAIN/load_functions')\n",
        "from download_from_gdrive import download_file_from_google_drive\n",
        "best_pretrained = \"1EYgrdsaffVJuQlcBSK4dGqO2BXBkSlZf\"\n",
        "if not os.path.exists(\"/content/DAIN/model_weights\"):\n",
        "  os.mkdir(\"/content/DAIN/model_weights\")\n",
        "file_path = \"/content/DAIN/model_weights/best.pth\"\n",
        "download_file_from_google_drive(best_pretrained, file_path)\n",
        "\n",
        "### Create folder for demonstration of for continue training with pretrained network\n",
        "if not os.path.exists(\"/content/DAIN/pretraine_model\"):\n",
        "  os.mkdir(\"/content/DAIN/pretraine_model\")\n",
        "import shutil\n",
        "shutil.copy(\"/content/DAIN/model_weights/best.pth\", \"/content/DAIN/pretraine_model/best.pth\")\n",
        "\n",
        "\n",
        "if not os.path.exists(\"/content/DAIN/MiddleBurySet\"):\n",
        "  os.mkdir(\"/content/DAIN/MiddleBurySet\")\n",
        "\n",
        "\n",
        "# Download simulated particle files:\n",
        "# Download demo file\n",
        "path_example_GT = \"/content/DAIN/demo/sim_particels_GT_512\"\n",
        "path_example_2DF_256 = \"/content/DAIN/demo/sim_particles_2DF_256\"\n",
        "path_example_4DF_256 = \"/content/DAIN/demo/sim_particles_2DF_2DF_256\"\n",
        "path_example_2DF_512 = \"/content/DAIN/demo/sim_particles_2DF_512\"\n",
        "path_example_4DF_512 = \"/content/DAIN/demo/sim_particles_2DF_2DF_512\"\n",
        "\n",
        "if not os.path.exists(\"/content/DAIN/demo\"):\n",
        "  os.mkdir(\"/content/DAIN/demo\")\n",
        "if not os.path.exists(path_example_GT):\n",
        "  os.mkdir(path_example_GT)\n",
        "if not os.path.exists(path_example_2DF_256):\n",
        "  os.mkdir(path_example_2DF_256)\n",
        "if not os.path.exists(path_example_4DF_256):\n",
        "  os.mkdir(path_example_4DF_256)\n",
        "if not os.path.exists(path_example_2DF_512):\n",
        "  os.mkdir(path_example_2DF_512)\n",
        "if not os.path.exists(path_example_4DF_512):\n",
        "  os.mkdir(path_example_4DF_512)\n",
        "\n",
        "\n",
        "example_GT = path_example_GT + \"/sim_particles_GT_512.tif\"\n",
        "example_2DF_256 = path_example_2DF_256 + \"/sim_particles_2DF_256.tif\"\n",
        "example_4DF_256 = path_example_4DF_256 + \"/sim_particle_2DF_2DF_256.tif\"\n",
        "example_2DF_512 = path_example_2DF_512 + \"/sim_particles_2DF_512.tif\"\n",
        "example_4DF_512 = path_example_4DF_512 + \"/sim_particles_2DF_2DF_512.tif\"\n",
        "\n",
        "demo_file_GT = \"1JQC78E0QBDbhiyNK-hPTYiiDWRkwfGd1\"\n",
        "demo_file_2DF_256 = \"1gKtp6y0aoYwMEefzC2oMJbXzfi3c_cCr\"\n",
        "demo_file_4DF_256 = \"1B1sls8n6irmjKkGNQculIe3M1ZeMEx2_\"\n",
        "demo_file_2DF_512 = \"1OPLXzCbTAviRVNUPYh7r67AKRYgSHmHn\"\n",
        "demo_file_4DF_512 = \"1CnMD97J-N9uP1letHwx7nMWgVCgroAdy\"\n",
        "\n",
        "\n",
        "download_file_from_google_drive(demo_file_GT, example_GT)\n",
        "download_file_from_google_drive(demo_file_2DF_256, example_2DF_256)\n",
        "download_file_from_google_drive(demo_file_4DF_256, example_4DF_256)\n",
        "download_file_from_google_drive(demo_file_2DF_512, example_2DF_512)\n",
        "download_file_from_google_drive(demo_file_4DF_512, example_4DF_512)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "QOEOuUY1YzG6"
      },
      "source": [
        "#@title ### Run cell to avoid Google Colab disconnect while waiting\n",
        "import time\n",
        "time.sleep(4800)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IWFtMREDS_LB"
      },
      "source": [
        "##Download additional demo files (optional)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PISP1qodS-UN",
        "cellView": "form"
      },
      "source": [
        "#@title ### Download additional demo files to play around with. After running this cell they will appear in the `/content/DAIN/demo` folder.\n",
        "# Download demo file\n",
        "import sys\n",
        "import os\n",
        "_ = (sys.path.append(\"/usr/local/lib/python3.6/site-packages\"))\n",
        "sys.path.insert(0,'/content/DAIN/load_functions')\n",
        "from download_from_gdrive import download_file_from_google_drive\n",
        "path_example_256 = \"/content/DAIN/demo/example_1_256_5D\"\n",
        "path_example_128 = \"/content/DAIN/demo/example_1_128_5D\"\n",
        "\n",
        "if not os.path.exists(\"/content/DAIN/demo\"):\n",
        "  os.mkdir(\"/content/DAIN/demo\")\n",
        "if not os.path.exists(path_example_256):\n",
        "  os.mkdir(path_example_256)\n",
        "if not os.path.exists(path_example_128):\n",
        "  os.mkdir(path_example_128)\n",
        "\n",
        "example_file_256 = path_example_256 + \"/example_256.tif\"\n",
        "example_file_128 = path_example_128 + \"/example_128.tif\"\n",
        "demo_file_256 = \"1Wx_c05Z3IAZlCz5ybhiEO8zPWtr5BzGR\"\n",
        "demo_file_128 = \"1ECeMd3J4R30nKmDijWqYuPrWk0Cq_FIa\"\n",
        "download_file_from_google_drive(demo_file_256, example_file_256)\n",
        "download_file_from_google_drive(demo_file_128, example_file_128)\n",
        "\n",
        "\n",
        "os.mkdir(\"/content/DAIN/demo/example_2_2048_3D\")\n",
        "os.mkdir(\"/content/DAIN/demo/example_2_512_4D\")\n",
        "os.mkdir(\"/content/DAIN/demo/example_2_256_4D\")\n",
        "os.mkdir(\"/content/DAIN/demo/example_2_128_4D\")\n",
        "\n",
        "link_file = \"1CgbYzFXI8fnF5sFuuic6r6IzwOg4yqPi\"\n",
        "demo_flie_location = \"/content/DAIN/demo/example_2_2048_3D/SHSY5Y_2048_3D.tif\"\n",
        "download_file_from_google_drive(link_file, demo_flie_location)\n",
        "\n",
        "file_512 = \"1FZ3G9dzDtsDGNPXe7sKDANuas4Tv_NI3\"\n",
        "demo_flie_location_512 = \"/content/DAIN/demo/example_2_512_4D/SHSY5Y_512_4D.tif\"\n",
        "download_file_from_google_drive(file_512, demo_flie_location_512)\n",
        "\n",
        "file_256 = \"1mxR17ojfazlTAe4HfqrzT_bX0X_kwNwx\"\n",
        "demo_flie_location_256 = \"/content/DAIN/demo/example_2_256_4D/SHSY5Y_256_4D.tif\"\n",
        "download_file_from_google_drive(file_256, demo_flie_location_256)\n",
        "\n",
        "file_128 = \"1_xMqFUkPq3J29i7tfTC6KeHvKVTB5mjL\"\n",
        "demo_flie_location_128 = \"/content/DAIN/demo/example_2_128_4D/SHSY5Y_128_4D.tif\"\n",
        "download_file_from_google_drive(file_128, demo_flie_location_128)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WT2RF0xWVWJm"
      },
      "source": [
        "\n",
        "# **3. Image manipulation tools** (Optional - help to prepare your workflow)\n",
        "#####Not necessary for training part\n",
        "---\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "R0AzGsX_6XWH"
      },
      "source": [
        "If you need to manipulate the images before performing training or testing you can do that here.\n",
        "Options:\n",
        "- Check the dimensions of a TIF file\n",
        "- Downscale image frequency in t-dimension\n",
        "- Downscale image frequency in z-dimension "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "N91WfL8TZaSe",
        "cellView": "form"
      },
      "source": [
        "#@title **Checking the dimensions of a selected TIF file**\n",
        "#@markdown Select a tif file to display the image dimensions.\n",
        "#@markdown The dimensions of your images should be one of the following:\n",
        "#@markdown ##**TZXYC** | **TZXY** | **TXY** | **ZXY**\n",
        "from skimage import io\n",
        "file_path = \"/content/DAIN/demo/sim_particles_2DF_2DF_256/sim_particle_2DF_2DF_256.tif\"#@param {type:\"string\"}\n",
        "\n",
        "img = io.imread(file_path)\n",
        "print(f\"The TIF image dimensions are {str(img.shape)}\")\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mM4x8B1tB3BB",
        "cellView": "form"
      },
      "source": [
        "#@markdown ##Downscale image frequency\n",
        "#@markdown The `downscale_t` box lets you reduce the image number in temporal dimension.\n",
        "#@markdown The `downscale_t` box lets you reduce the image number in axial dimension.\n",
        "#@markdown The `source_path` field asks for the folder which contain your image sequences.\n",
        "\n",
        "\n",
        "downscale_t = True#@param{type:\"boolean\"}\n",
        "downscale_z = False#@param{type:\"boolean\"}\n",
        "#@markdown This options `limit_t_dim` allows the user limit the resulting image sequence to a selected `target_frame_number` (e.g., if `cut_target_frame_number` = 10 with `downscale_factor` = 2 and input dimensions = 25:512:512 in an TXY file, it will result in an image sequence of 10:512:512 dimensions)\n",
        "limit_t_dim = False#@param{type:\"boolean\"}\n",
        "\n",
        "cut_target_frame_number = 5#@param{type:\"number\"}\n",
        "\n",
        "def load_img(img_path):\n",
        "    img = io.imread(img_path)\n",
        "    img, use_RGB = correct_channels(img)\n",
        "    if img.shape[-1]==3:\n",
        "      use_RGB = True\n",
        "      t, z, y_dim, x_dim, _ = img.shape \n",
        "      print(\"This image will be processed as a RGB image\")\n",
        "    else:\n",
        "      use_RGB = False\n",
        "      t, z, y_dim, x_dim = img.shape \n",
        "    print(\"The image dimensions are: \" + str(img.shape))\n",
        "    return t, z, y_dim,x_dim, img, use_RGB\n",
        "\n",
        "\n",
        "def correct_channels(img):\n",
        "  '''For 2D + T (with or without RGB) a artificial z channel gets created'''\n",
        "  if img.shape[-1] ==3:\n",
        "    use_RGB = True\n",
        "  else:\n",
        "    use_RGB = False\n",
        "  if len(img.shape) ==4 and use_RGB:\n",
        "    t, x, y, c = img.shape\n",
        "    zeros = np.zeros((t,1,y,x,c),dtype=np.uint8)\n",
        "    zeros[:,0,:,:,:] = img\n",
        "    img = zeros\n",
        "  elif len(img.shape) ==3 and not use_RGB:\n",
        "    t, x, y = img.shape\n",
        "    zeros = np.zeros((t,1,y,x),dtype=np.uint8)\n",
        "    zeros[:,0,:,:] = img\n",
        "    img = zeros\n",
        "  return img, use_RGB\n",
        "\n",
        "#@title Downsampling of image\n",
        "# TODO- fix for RGB\n",
        "\n",
        "from skimage import io\n",
        "import numpy as np\n",
        "import os\n",
        "\n",
        "downscale_factor =  2#@param {type:\"number\"}\n",
        "source_path = \"/content/DAIN/demo/sim_particles_2DF_256\"#@param {type:\"string\"}\n",
        "source_name = os.path.basename(source_path)\n",
        "save_parent = os.path.dirname(source_path)\n",
        "save_folder = source_name + f\"_{downscale_factor}DF\"\n",
        "save_path = os.path.join(save_parent, save_folder)\n",
        "if not os.path.exists(save_path):\n",
        "  os.mkdir(save_path)\n",
        "\n",
        "\n",
        "\n",
        "flist = os.listdir(source_path)\n",
        "flist = [f for f in flist if f.endswith(\"tif\")]\n",
        "\n",
        "for name in flist:\n",
        "  file_path = os.path.join(source_path, name)\n",
        "  t, z, y_dim,x_dim, img, use_RGB = load_img(file_path)\n",
        "\n",
        "\n",
        "  if use_RGB == False:\n",
        "    if downscale_z and downscale_t:\n",
        "      label = \"_t_z_\"\n",
        "      img_new = img[::downscale_factor,::downscale_factor,:,:]\n",
        "    elif downscale_t:\n",
        "      label = \"_t_\"\n",
        "      img_new = img[::downscale_factor,:,:,:]\n",
        "    elif downscale_z:\n",
        "      label = \"_z_\"\n",
        "      img_new = img[:,::downscale_factor,:,:]\n",
        "  else:\n",
        "    if downscale_z and downscale_t:\n",
        "     label = \"_t_z_\"\n",
        "     img_new = img[::downscale_factor,::downscale_factor,:,:,:]\n",
        "    elif downscale_t:\n",
        "      label = \"_t_\"\n",
        "      img_new = img[::downscale_factor,:,:,:]\n",
        "    elif downscale_z:\n",
        "      label = \"_z_\"\n",
        "      img_new = img[:,::downscale_factor,:,:]\n",
        "\n",
        "  save_name = f\"DF_{downscale_factor}{label}\" +  name\n",
        "  save_path_file = os.path.join(save_path, save_name)\n",
        "  if use_RGB == False:\n",
        "    if limit_t_dim:\n",
        "      img_new = img_new[:cut_target_frame_number,:,:,:]\n",
        "  if use_RGB == True:\n",
        "    if limit_t_dim:\n",
        "      img_new = img_new[:cut_target_frame_number,:,:,:,:]\n",
        "  print(img_new.shape)\n",
        "  io.imsave(save_path_file, img_new)\n",
        "\n",
        "  # save_name = f\"DF2_DF_{downscale_factor}{label}\" +  name\n",
        "  # save_path_file = os.path.join(save_path, save_name)\n",
        "  # img_DS = img_new[::downscale_factor,::downscale_factor,:,:]\n",
        "  # print(img_DS.shape)\n",
        "\n",
        "  # io.imsave(save_path_file, img_DS)\n",
        "\n",
        "  # img_temp = convert(img_temp, 0, 255, np.uint8)\n",
        "  # io.imsave(\"DS-{}-{}.tif\".format(downscale_factor, name),img_new)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d_Ov3NRSJ-gw"
      },
      "source": [
        "\n",
        "# **4. Training**\n",
        "---"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MX4N4Md6bHZF"
      },
      "source": [
        "###**4.1 Prepare Data for Training**\n",
        "If you have small datasets for finetuning the network you can run the cells in this section to prepare the data. If you want to prepare bigger dataset (more than 5 GB) we recommend you to use the python scrips on your computer using python provided in the following folder: `/content/DAIN/scripts`\n",
        "You can upload the created folders to Google Colab via Google Drive or Google Cloud Storage and load it in the training stage. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hboFxmZLKDS4",
        "cellView": "form"
      },
      "source": [
        "#@title 4.1.2 Data preparation for training the network\n",
        "\n",
        "import sys\n",
        "import os\n",
        "import csv\n",
        "from tqdm import tqdm\n",
        "from skimage import io\n",
        "from datetime import datetime\n",
        "\n",
        "_ = (sys.path.append(\"/usr/local/lib/python3.6/site-packages\"))\n",
        "sys.path.insert(0,'/content/DAIN/load_functions')\n",
        "from prepare_split_images import make_folder_with_date\n",
        "from prepare_split_images import bcolors\n",
        "from manipulate_listdatasets import manipulate_listdatasets\n",
        "from prepare_dataset_train_test_folders import load_img\n",
        "from prepare_dataset_train_test_folders import correct_channels\n",
        "from prepare_split_images import split_img_small\n",
        "\n",
        "# Splitting of the image\n",
        "Source_path = \"/content/DAIN/demo/sim_particles_2DF_512\" #@param {type:\"string\"}\n",
        "Parent_path = \"/\".join(Source_path.split(\"/\")[:-1])\n",
        "\n",
        "img_size = 512 #@param [\"64\", \"128\", \"256\", \"512\", \"1024\", \"2048\"] {type:\"raw\"}\n",
        "if img_size<=512:\n",
        "  divisor = img_size\n",
        "else: \n",
        "  divisor = 512\n",
        "\n",
        "\n",
        "\n",
        "# Write a log file for keeping track of the names\n",
        "log_path_file = Parent_path + \"/\" + \"split_log.csv\"\n",
        "\n",
        "with open(log_path_file, 'w', newline='') as file:\n",
        "      writer = csv.writer(file)\n",
        "      writer.writerow([\"file_name\", \"split_name\"])\n",
        "\n",
        "# Get all the paths of the images\n",
        "img_list = [f for f in os.listdir(Source_path) if f.endswith('.tif')]\n",
        "\n",
        "# Create parent-folder where each split execution will be saved\n",
        "aug_saving_path = Parent_path+'/spit_source_DAIN'\n",
        "if not os.path.exists(aug_saving_path):\n",
        "  os.mkdir(aug_saving_path)\n",
        "\n",
        "\n",
        "# changes the dimensions for the listdatasets.py file scriptn - might not be needed\n",
        "manipulate_listdatasets(divisor)\n",
        "\n",
        "# create a folder where split images are being stored\n",
        "# dirname_path = os.path.dirname((Source_path))\n",
        "split_img_folder_path = make_folder_with_date(aug_saving_path, \"split\")\n",
        "\n",
        "split_img_small(img_list, Source_path, divisor, split_img_folder_path, log_path_file)\n",
        "\n",
        "print(bcolors.WARNING + \"Finished section 1: split_img_small \")\n",
        "\n",
        "\n",
        "\n",
        "##############################################\n",
        "#@title Prepare data for Training or Interpolation\n",
        "\n",
        "from prepare_dataset_train_test_folders import data_train_test_preparation\n",
        "\n",
        "# create new folder\n",
        "save_location = \"/\".join(split_img_folder_path.split(\"/\")[:-1])\n",
        "folder_name = f\"{divisor}_img_separation\"\n",
        "save_location = os.path.join(save_location,folder_name)\n",
        "if not os.path.exists(save_location):\n",
        "  os.mkdir(save_location)\n",
        "os.chdir(save_location)  \n",
        "\n",
        "#param [\"prep_t_train\", \"prep_z_train\", \"prep_predict_t\", \"prep_predict_z\", \"upsample_t\", \"upsample_z\"]\n",
        "folder_option = \"prep_t_train\" \n",
        "split_training_test = 0.1\n",
        "\n",
        "# split_folder, folder_steps, folder_gt, train_folder, sub_save_location\n",
        "_, _, _, _, sub_save_location = data_train_test_preparation(folder_option, split_img_folder_path, save_location, split_training_test)\n",
        "print(bcolors.WARNING + \"Finished section 2: data train test preparation\")\n",
        "\n",
        "###############################################\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wEmV2hTZbT-z"
      },
      "source": [
        "### **4.2 Perform fine-tuning of pretrained network**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "97wQfLN1bPwS",
        "cellView": "form"
      },
      "source": [
        "#@title 4.2.1 Select data and parameters for training\n",
        "\n",
        "import sys\n",
        "sys.path.insert(0,'/content/DAIN/load_functions')\n",
        "import shutil\n",
        "import os\n",
        "import time\n",
        "from prepare_split_images import make_folder_with_date\n",
        "# because of: scipy  ImportError: cannot import name 'imread'\n",
        "!yes y | pip uninstall scipy\n",
        "!pip install scipy==1.1.0\n",
        "\n",
        "#1. Select parameters ##@param [\"t-dimension\", \"z-dimension\"]\n",
        "# select_train_dimension = \"t-dimension\" \n",
        "#@markdown #### If you don't select your own pretrained model the `best.pth` model in the `model_weights` folder will be fine tuned.\n",
        "select_your_pretrained_model = False #@param {type:\"boolean\"}\n",
        "pretrained_model_location =  \"\"#@param {type:\"string\"}\n",
        "#@markdown #### Select this if you already prepared your training data previously and you did not run the cell above. The correct folder to select is called date_`prep_t_train`.\n",
        "manually_select_train_folder = False #@param {type:\"boolean\"}\n",
        "if manually_select_train_folder:\n",
        "  sub_save_location =  \"/content/DAIN/demo/spit_source_DAIN/512_img_separation/20211018_094552_prep_t_train\"#@param {type:\"string\"}\n",
        "\n",
        "#@markdown ###Choose Parameters\n",
        "\n",
        "epochs =  1#@param {type:\"number\"}\n",
        "batch_size =  1#@param {type:\"number\"}\n",
        "lr = 0.0005#@param {type:\"number\"}\n",
        "\n",
        "\n",
        "#@markdown #### Select a folder on your google drive where the fine-tuned model should be saved after training\n",
        "model_backup_location =  \"/content/sample_data\"#@param {type:\"string\"}\n",
        "model_backup_sublocation = model_backup_location+'/DAIN_model_archive'\n",
        "if not os.path.exists(model_backup_sublocation):\n",
        "  os.mkdir(model_backup_sublocation)\n",
        "DAIN_model_backup = make_folder_with_date(model_backup_sublocation, \"DAIN_model_backup\")\n",
        "\n",
        "load_data_from_folder =  sub_save_location\n",
        "# select_folder = True #@param {type:\"boolean\"}\n",
        "\n",
        "\n",
        "os.chdir(\"/content/DAIN\")\n",
        "#2. Insert the above values into train_model.sh\n",
        "!chmod u+x train_model.sh\n",
        "!sed -i \"s/NUMEPOCH=.*/NUMEPOCH=$epochs/g\" train_model.sh #change the number of training iterations (steps)\n",
        "!sed -i \"s/BATCHSIZE=.*/BATCHSIZE=$batch_size/g\" train_model.sh #change the number of training iterations (steps)\n",
        "!sed -i \"s|DATASETPATH=.*|DATASETPATH='$load_data_from_folder'|g\" train_model.sh #change the number of training iterations (steps)\n",
        "!sed -i \"s|LR=.*|LR=$lr|g\" train_model.sh #change the number of training iterations (steps)\n",
        "\n",
        "!chmod u+x continue_train_model.sh\n",
        "!sed -i \"s/NUMEPOCH=.*/NUMEPOCH=$epochs/g\" continue_train_model.sh #change the number of training iterations (steps)\n",
        "!sed -i \"s/BATCHSIZE=.*/BATCHSIZE=$batch_size/g\" continue_train_model.sh #change the number of training iterations (steps)\n",
        "!sed -i \"s|DATASETPATH=.*|DATASETPATH='$load_data_from_folder'|g\" continue_train_model.sh #change the number of training iterations (steps)\n",
        "!sed -i \"s|PRETRAINED=.*|PRETRAINED='$model_folder_name'|g\" continue_train_model.sh #change the number of training iterations (steps)\n",
        "!sed -i \"s|LR=.*|LR=$lr|g\" continue_train_model.sh #change the number of training iterations (steps)\n",
        "# because sometimes it takes time until linux and colab syncronize their file systems from the folder preparation\n",
        "time.sleep(5)\n",
        "\n",
        "if select_your_pretrained_model == False:\n",
        "  !cp -r \"/content/DAIN/model_weights/best.pth\" \"/content/best.pth\"\n",
        "  !rm -r \"/content/DAIN/model_weights\"\n",
        "  !mkdir \"/content/DAIN/model_weights\"\n",
        "  !cp -r \"/content/best.pth\" \"/content/DAIN/model_weights\" \n",
        "  !rm -r \"/content/best.pth\"\n",
        "  #Training parameters in DAIN are indicated in the train_model.sh file.\n",
        "  #Here, we edit this file to include the desired parameters\n",
        "  \n",
        "  # # because sometimes it takes time until linux and colab syncronize their file systems from the folder preparation\n",
        "  # time.sleep(5)\n",
        "  !/content/DAIN/train_model.sh \n",
        "\n",
        "  #Save model on google drive\n",
        "  shutil.rmtree(DAIN_model_backup)\n",
        "  shutil.copytree(\"/content/DAIN/model_weights\", DAIN_model_backup)\n",
        "  print(\"Trained model was saved in the following folder: \\n%s\"%DAIN_model_backup)\n",
        "\n",
        "elif select_your_pretrained_model == True:\n",
        "  !rm -r \"/content/DAIN/model_weights\"\n",
        "  !mkdir \"/content/DAIN/model_weights\"\n",
        "\n",
        "  model_folder_name = pretrained_model_location.split(\"/\")[-1]\n",
        "  model_folder_path = \"/content/DAIN/model_weights/%s\"%model_folder_name\n",
        "  shutil.copytree(pretrained_model_location, model_folder_path)\n",
        "  shutil.copy(model_folder_path+\"/best.pth\",\"/content/DAIN/model_weights/best.pth\")\n",
        "    \n",
        "  # number_of_images =  len(os.listdir(Training_source))\n",
        " \n",
        "  os.chdir(\"/content/DAIN\")\n",
        "  !/content/DAIN/continue_train_model.sh \n",
        "\n",
        "  #Save model on google drive\n",
        "  shutil.rmtree(DAIN_model_backup)\n",
        "\n",
        "  shutil.copytree(\"/content/DAIN/model_weights\", DAIN_model_backup)\n",
        "  print(\"Trained model was saved in the following folder: \\n%s\"%DAIN_model_backup)\n",
        "\n",
        "  #@markdown ####**Note:** Training might take several hours. If your google colab disconnects you can continue your training by using the same data and saved model weights in **section 2.1.9**\n",
        "\n",
        "#@markdown ###**Note:** existing models in the following folder: `/content/DAIN/model_weights` will be deleted!!!\n",
        "#@markdown #### Save model first before starting new training!! "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9D4pLOqIGAJc",
        "cellView": "form"
      },
      "source": [
        "# @markdown ###4.2.3 Play the cell to show figure of training errors\n",
        "import os\n",
        "import csv\n",
        "from matplotlib import pyplot as plt\n",
        "\n",
        "#@markdown ####If you didn't continue training but just want to see the result of a saved trained model enter the path of the folder and run the code.\n",
        "#@markdown ####The latest model is saved either on your selected google drive folder and is still present in the following folder with your date of execution `/content/DAIN/model_weights/65228-Thu-Jun-04-13:23`\n",
        "check_model = \"/content/DAIN/model_weights/4272-Wed-Nov-03-06-05\"#@param {type:\"string\"}\n",
        "# Check the trainings-progress by checking the loss of the network\n",
        "# root_path = \"/content/DAIN/model_weights\"\n",
        "# if os.path.exists(DAIN_model_backup):\n",
        "#   model_name = DAIN_model_backup   ### no \"/\" just the name\n",
        "# else:\n",
        "#   model_name = check_model\n",
        "file_name = \"log.txt\"\n",
        "\n",
        "log_file_path = os.path.join(check_model,file_name)\n",
        "x = []\n",
        "y1 = []\n",
        "\n",
        "print(os.path.exists(log_file_path))\n",
        "with open(log_file_path,'r') as csvfile:\n",
        "    plots = csv.reader(csvfile, delimiter=',')\n",
        "    # print(plots)\n",
        "    next(plots)\n",
        "    for row in plots:\n",
        "        x.append(float(row[0]))\n",
        "        y1.append(float(row[5]))\n",
        "        \n",
        "plt.plot(x,y1)\n",
        "plt.title('Model loss')\n",
        "plt.ylabel('Loss')\n",
        "plt.xlabel('Iteration')\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2Hq0tP5bXWCS"
      },
      "source": [
        "\n",
        "# **5. Perform DAIN Frame Interpolation**\n",
        "---"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "W8YK7Ed7z_UG",
        "cellView": "form"
      },
      "source": [
        "############ SELECT OPTIONS ################\n",
        "#########Split images in single#############\n",
        "############################################\n",
        "import sys\n",
        "_ = (sys.path.append(\"/usr/local/lib/python3.6/site-packages\"))\n",
        "sys.path.insert(0,'/content/DAIN/load_functions')\n",
        "import os\n",
        "from tqdm import tqdm\n",
        "from datetime import datetime\n",
        "from skimage import io\n",
        "import csv\n",
        "import time\n",
        "from prepare_split_images import make_folder_with_date\n",
        "from prepare_split_images import bcolors\n",
        "from manipulate_listdatasets import manipulate_listdatasets\n",
        "from prepare_split_images import split_img_small\n",
        "from prepare_split_images import bcolors\n",
        "\n",
        "\n",
        "#@markdown Make sure that all the images in the folder have all **the SAME dimensions**\n",
        "\n",
        "# Splitting of the image\n",
        "Source_path = \"/content/DAIN/demo/sim_particles_2DF_2DF_512\" #@param {type:\"string\"}\n",
        "Saving_path = \"/\".join(Source_path.split(\"/\")[:-1])\n",
        "\n",
        "\n",
        "#@markdown Provide the size of the images\n",
        "img_size = 512 #@param [\"64\", \"128\", \"256\", \"512\", \"1024\", \"2048\"] {type:\"raw\"}\n",
        "if img_size<=512:\n",
        "  divisor = img_size\n",
        "else: \n",
        "  divisor = 512\n",
        "\n",
        "perform_t_interpolation = True #@param {type:\"boolean\"}\n",
        "perform_z_interpolation = False #@param {type:\"boolean\"}\n",
        "\n",
        "if perform_t_interpolation == True:\n",
        "  folder_option_1 = \"upsample_t\"\n",
        "  if perform_z_interpolation == True:\n",
        "    folder_option_2 =\"upsample_z\"\n",
        "if perform_z_interpolation == True and perform_t_interpolation == False:\n",
        "  folder_option_1 = \"upsample_z\"\n",
        "\n",
        "\n",
        "if perform_z_interpolation == False and perform_t_interpolation == False:\n",
        "   print(bcolors.WARNING + f\"No image improvement selected\")\n",
        "   exit()\n",
        "\n",
        "# Write a log file for keeping track of the names\n",
        "log_path_file_1 = Saving_path + \"/\" + \"split_log_1.csv\"\n",
        "\n",
        "with open(log_path_file_1, 'w', newline='') as file:\n",
        "      writer = csv.writer(file)\n",
        "      writer.writerow([\"file_name\", \"split_name\"])\n",
        "\n",
        "# Create parent-folder where each split execution will be saved\n",
        "aug_saving_path = Saving_path+'/spit_source_DAIN'\n",
        "if not os.path.exists(aug_saving_path):\n",
        "  os.mkdir(aug_saving_path)\n",
        "\n",
        "# Get all the paths of the images\n",
        "img_list = [f for f in os.listdir(Source_path) if f.endswith('.tif')]\n",
        "\n",
        "# changes the dimensions for the listdatasets.py file scriptn - might not be needed\n",
        "manipulate_listdatasets(divisor)\n",
        "\n",
        "# create a folder where split images are being stored\n",
        "# dirname_path = os.path.dirname((Source_path))\n",
        "split_img_folder_path = make_folder_with_date(aug_saving_path, \"split\")\n",
        "\n",
        "use_RGB = split_img_small(img_list, Source_path, divisor, split_img_folder_path, log_path_file_1)\n",
        "\n",
        "############ PREP DAIN DATA ################\n",
        "#########Get right solder system############\n",
        "############################################\n",
        "from prepare_dataset_train_test_folders import data_train_test_preparation\n",
        "# create new folder\n",
        "save_location = \"/\".join(split_img_folder_path.split(\"/\")[:-1])\n",
        "folder_name = f\"{divisor}_img_separation\"\n",
        "save_location = os.path.join(save_location,folder_name)\n",
        "if not os.path.exists(save_location):\n",
        "  os.mkdir(save_location)\n",
        "os.chdir(save_location)  \n",
        "\n",
        "\n",
        "# split_folder, folder_steps, folder_gt, train_folder_, sub_save_folder\n",
        "split_folder_1, _, _, _, _ = data_train_test_preparation(folder_option_1, split_img_folder_path, save_location, _)\n",
        "\n",
        "############ LOAD DATA IN REPO ##############\n",
        "#########Into MiddleBurySet folder###########\n",
        "#############################################\n",
        "import shutil\n",
        "!rm -r \"/content/DAIN/MiddleBurySet/other-data\"\n",
        "\n",
        "# load_folder = False #@param {type:\"boolean\"}\n",
        "\n",
        "# if load_folder == False:\n",
        "if (perform_t_interpolation == True or perform_z_interpolation == True):\n",
        "  split_folder_path = split_folder_1\n",
        "# elif  (folder_option == \"prep_predict_z\" or folder_option == \"prep_predict_t\"):\n",
        "#   split_folder_path = folder_steps\n",
        "# else:\n",
        "#   split_folder_path = \"/content/gdrive/MyDrive/1.2_BIG_DATA_PhD_Project_2/6.ZoomInterpolation/20210208_cell_data_8bit+CARE_DAIN/spit_augmented_source_DAIN/512_img_separation/2021021616_161414_upsample_z/steps\"#@param {type:\"string\"}\n",
        "\n",
        "#@markdown Here you can either select `use_provided_pretrained_model` and it will perform the interpolation with the provided pretrained model or if unchecked it will load your fine-tuned network from the provided folder below. The folder you should select should look like that `43012-Thu-Oct-14-11-51`. \n",
        "#Select the parent folder with the date and time as a source\n",
        "\n",
        "# copy the prepared test data into the file\n",
        "# folder = split_folder_path.split(\"/\")[-1] \n",
        "destination_path = \"/content/DAIN/MiddleBurySet\"\n",
        "\n",
        "if not os.path.exists(destination_path):\n",
        "  os.mkdir(destination_path)\n",
        "destination_folder = os.path.join(destination_path, \"other-data\")\n",
        "shutil.copytree(split_folder_path, destination_folder)\n",
        "os.chdir(destination_folder)\n",
        "if os.path.exists(\"name_log.txt\"):\n",
        "  os.remove(\"name_log.txt\")\n",
        "if os.path.exists(\".ipynb_checkpoints\"):\n",
        "  os.remove(\".ipynb_checkpoints\")\n",
        "\n",
        "\n",
        "######### LOAD PRETRAINED NETWORK ###########\n",
        "#########Choose or use provided one##########\n",
        "#############################################\n",
        "#@markdown ### **Load best trainingsmodel** \n",
        "use_provided_pretrained_model = True #@param {type:\"boolean\"}\n",
        "if use_provided_pretrained_model == True:\n",
        "  shutil.copy(\"/content/DAIN/model_weights/best.pth\", \"/content/best.pth\")\n",
        "  shutil.rmtree(\"/content/DAIN/model_weights\")\n",
        "  os.mkdir(\"/content/DAIN/model_weights\")\n",
        "  shutil.copy(\"/content/best.pth\", \"/content/DAIN/model_weights/best.pth\")\n",
        "  os.mkdir(\"/content/DAIN/model_weights/best\")\n",
        "  shutil.copy(\"/content/best.pth\", \"/content/DAIN/model_weights/best/best.pth\")\n",
        "  model_folder_name = \"best\"\n",
        "\n",
        "else:\n",
        "  #@markdown ###The model in the selected folder that is used must be named **\"best.pth\"**\n",
        "  best_model_path = \"/content/DAIN/model_weights/68932-Thu-Aug-26-08-23\"#@param {type:\"string\"}\n",
        "  model_folder_name = best_model_path.split(\"/\")[-1]\n",
        "  model_folder_path = \"/content/DAIN/model_weights/selected_%s\"%model_folder_name\n",
        "  if os.path.exists(model_folder_path):\n",
        "    shutil.rmtree(model_folder_path)\n",
        "  shutil.copytree(best_model_path, model_folder_path)\n",
        "  \n",
        "\n",
        "######### RUN FIRST INTERPOLATION ###########\n",
        "#########        T-Dimension       ##########\n",
        "#############################################\n",
        "!yes y | pip uninstall scipy\n",
        "!pip install scipy==1.1.0\n",
        "\n",
        "#1. Add permissions to execute_interpolation.sh\n",
        "os.chdir(\"/content/DAIN\")\n",
        "!chmod u+x execute_interpolation.sh\n",
        "\n",
        "#@markdown Select the number of interpolated images between two given ones\n",
        "interpolation_number = \"1\" #@param [\"1\", \"2\", \"3\"]\n",
        "if interpolation_number== \"1\":\n",
        "  timestep =  0.5\n",
        "elif interpolation_number== \"2\":\n",
        "  timestep =  0.33\n",
        "elif interpolation_number== \"3\":\n",
        "  timestep =  0.25\n",
        "\n",
        "# number_of_images =  len(os.listdir(Training_source))\n",
        "\n",
        "!sed -i \"s/MODEL=.*/MODEL='$model_folder_name'/g\" execute_interpolation.sh #change the number of training iterations (steps)\n",
        "!sed -i \"s/TIMESTEP=.*/TIMESTEP=$timestep/g\" execute_interpolation.sh #change the number of training iterations (steps)\n",
        "\n",
        "# because sometimes it takes time until linux and colab syncronize their file systems from the folder preparation\n",
        "time.sleep(1)\n",
        "#------------------------\n",
        "# #3. Insert the above values into train_model.sh\n",
        "!/content/DAIN/execute_interpolation.sh\n",
        "\n",
        "\n",
        "###### PREPARE RECONSTRUCTION FOLDER ########\n",
        "#########                          ##########\n",
        "#############################################\n",
        "####If not selected load folder it will take the latest interpolation results created in the following folder: /content/DAIN/MiddleBurySet/other-result-author\n",
        "from reconstruct_image import restructure_folder_for_processing\n",
        "\n",
        "folder_list = os.listdir(\"/content/DAIN/MiddleBurySet/other-result-author\")\n",
        "folder_list = [i for i in folder_list if not \"hdf5\" in i]\n",
        "folder_list.sort()\n",
        "interpolate_location = os.path.join(\"/content/DAIN/MiddleBurySet/other-result-author\", folder_list[-1])\n",
        "\n",
        "reprocessed_folder = restructure_folder_for_processing(interpolate_location, Saving_path, log_path_file_1, divisor, folder_option_1, use_RGB)\n",
        "\n",
        "###### SAVE FIRST INTERPOLATED IMAGE ########\n",
        "#########                          ##########\n",
        "#############################################\n",
        "\n",
        "# load_folder = False #@param {type:\"boolean\"}\n",
        "\n",
        "# if load_folder:\n",
        "#   interpolate_location = \"/content/DAIN/MiddleBurySet/processed/20210218_153839_experiment\"#@param {type:\"string\"}\n",
        "# else:\n",
        "interpolate_location = reprocessed_folder\n",
        "# create a list of the identifyer for \n",
        "\n",
        "from reconstruct_image import save_interpolated_image\n",
        "Source_path_2 = save_interpolated_image(interpolate_location, Saving_path, log_path_file_1, folder_option_1, divisor, use_RGB)\n",
        "print(bcolors.WARNING + f\"Finished first interpolation step\")\n",
        "\n",
        "remove_folder = [Source_path_2]\n",
        "\n",
        "#### PERFORM SECOND DIM INTERPOLATION #######\n",
        "##########       Z-Dimansion      ###########\n",
        "#############################################\n",
        "# second interpolation step for the z dimension\n",
        "if perform_t_interpolation == True and perform_z_interpolation == True:\n",
        "\n",
        "  Saving_path_2 = \"/\".join(Source_path_2.split(\"/\")[:-1])\n",
        "\n",
        "  # Write a log file for keeping track of the names\n",
        "  log_path_file_2 = Saving_path_2 + \"/\" + \"split_log_2.csv\"\n",
        "\n",
        "  with open(log_path_file_2, 'w', newline='') as file:\n",
        "        writer = csv.writer(file)\n",
        "        writer.writerow([\"file_name\", \"split_name\"])\n",
        "\n",
        "  # Create parent-folder where each split execution will be saved\n",
        "  aug_saving_path = Saving_path+'/spit_source_DAIN'\n",
        "  if not os.path.exists(aug_saving_path):\n",
        "    os.mkdir(aug_saving_path)\n",
        "\n",
        "  # Get all the paths of the images\n",
        "  img_list = [f for f in os.listdir(Source_path_2) if f.endswith('.tif')]\n",
        "\n",
        "  # create a folder where split images are being stored\n",
        "  split_img_folder_path = make_folder_with_date(aug_saving_path, \"split\")\n",
        "\n",
        "  use_RGB = split_img_small(img_list, Source_path_2, divisor, split_img_folder_path, log_path_file_2)\n",
        "\n",
        "\n",
        "  ############ PREP DAIN DATA ################\n",
        "  #########Get right solder system############\n",
        "  ##############Z-Dimension###################\n",
        "  from prepare_dataset_train_test_folders import data_train_test_preparation\n",
        "  # create new folder\n",
        "  save_location = \"/\".join(split_img_folder_path.split(\"/\")[:-1])\n",
        "  folder_name = f\"{divisor}_img_separation\"\n",
        "  save_location = os.path.join(save_location,folder_name)\n",
        "  if not os.path.exists(save_location):\n",
        "    os.mkdir(save_location)\n",
        "  os.chdir(save_location)  \n",
        "\n",
        "\n",
        "  # split_folder, folder_steps, folder_gt, train_folder\n",
        "  split_folder_2, _, _, _, _ = data_train_test_preparation(folder_option_2, split_img_folder_path, save_location, _)\n",
        "\n",
        "\n",
        "\n",
        "  ############ LOAD DATA IN REPO ##############\n",
        "  #########Into MiddleBurySet folder###########\n",
        "  #############Z-Dimension##################\n",
        "\n",
        "  #Load downsampled testdata into DAIN folder on Colab machine\n",
        "  import shutil\n",
        "  !rm -r \"/content/DAIN/MiddleBurySet/other-data\"\n",
        "\n",
        "  # load_folder = False #@param {type:\"boolean\"}\n",
        "\n",
        "  # if load_folder == False:\n",
        "  if (perform_t_interpolation == True or perform_z_interpolation == True):\n",
        "    split_folder_path = split_folder_2\n",
        "\n",
        "  destination_path = \"/content/DAIN/MiddleBurySet\"\n",
        "\n",
        "  if not os.path.exists(destination_path):\n",
        "    os.mkdir(destination_path)\n",
        "  destination_folder = os.path.join(destination_path, \"other-data\")\n",
        "  shutil.copytree(split_folder_path, destination_folder)\n",
        "  os.chdir(destination_folder)\n",
        "  if os.path.exists(\"name_log.txt\"):\n",
        "    os.remove(\"name_log.txt\")\n",
        "  if os.path.exists(\".ipynb_checkpoints\"):\n",
        "    os.remove(\".ipynb_checkpoints\")\n",
        "\n",
        "\n",
        "  ######### RUN FIRST INTERPOLATION ###########\n",
        "  #########        T-Dimension       ##########\n",
        "  #############################################\n",
        "  !yes y | pip uninstall scipy\n",
        "  !pip install scipy==1.1.0\n",
        "\n",
        "  #1. Add permissions to execute_interpolation.sh\n",
        "  os.chdir(\"/content/DAIN\")\n",
        "  !chmod u+x execute_interpolation.sh\n",
        "\n",
        "  !sed -i \"s/MODEL=.*/MODEL='$model_folder_name'/g\" execute_interpolation.sh #change the number of training iterations (steps)\n",
        "  !sed -i \"s/TIMESTEP=.*/TIMESTEP=$timestep/g\" execute_interpolation.sh #change the number of training iterations (steps)\n",
        "\n",
        "  # because sometimes it takes time until linux and colab syncronize their file systems from the folder preparation\n",
        "  time.sleep(2)\n",
        "  #------------------------\n",
        "  # #3. Insert the above values into train_model.sh\n",
        "  !/content/DAIN/execute_interpolation.sh\n",
        "\n",
        "  print(\"Finished z-interpolation step\")\n",
        "\n",
        "  ###### PREPARE RECONSTRUCTION FOLDER ########\n",
        "  #########                          ##########\n",
        "  ###############Z-Dimension###################\n",
        "  from reconstruct_image import restructure_folder_for_processing\n",
        "\n",
        "  # folder_option = \"upsample_t\"\n",
        "  # load_folder = False #@param {type:\"boolean\"}\n",
        "  # if load_folder:\n",
        "  #   interpolate_location = \"/content/DAIN/MiddleBurySet/other-result-author/20210216_162649\"#@param {type:\"string\"}\n",
        "  # else:\n",
        "  folder_list = os.listdir(\"/content/DAIN/MiddleBurySet/other-result-author\")\n",
        "  folder_list = [i for i in folder_list if not \"hdf5\" in i]\n",
        "  folder_list.sort()\n",
        "  interpolate_location = os.path.join(\"/content/DAIN/MiddleBurySet/other-result-author\", folder_list[-1])\n",
        "\n",
        "  reprocessed_folder = restructure_folder_for_processing(interpolate_location, Saving_path, log_path_file_2, divisor, folder_option_2, use_RGB)\n",
        "\n",
        "\n",
        "  ###### SAVE FIRST INTERPOLATED IMAGE ########\n",
        "  #########                          ##########\n",
        "  #############################################\n",
        "  #Save interpolated image**\n",
        "\n",
        "  from reconstruct_image import save_interpolated_image\n",
        "\n",
        "  #load_folder = False #@param {type:\"boolean\"}\n",
        "  #if load_folder:\n",
        "  #  interpolate_location = \"/content/DAIN/MiddleBurySet/processed/20210218_153839_experiment\"#@param {type:\"string\"}\n",
        "  #else:\n",
        "  interpolate_location = reprocessed_folder\n",
        "  # create a list of the identifyer for \n",
        "  Source_path_3 = save_interpolated_image(interpolate_location, Saving_path, log_path_file_2, folder_option_2, divisor, use_RGB)\n",
        "  remove_folder.append(Source_path_3)\n",
        "\n",
        "  print(bcolors.WARNING + f\"Finished second interpolation step\")\n",
        "\n",
        "\n",
        "# saves final result in the folder named \"final result\"\n",
        "final = Saving_path +\"/final_result\"\n",
        "\n",
        "try:\n",
        "  if os.path.exists(Source_path_3):\n",
        "    if os.path.exists(final):\n",
        "        shutil.rmtree(final)\n",
        "    shutil.copytree(Source_path_3, final)\n",
        "except:\n",
        "  if os.path.exists(Source_path_2):\n",
        "    if os.path.exists(final):\n",
        "      shutil.rmtree(final)\n",
        "    shutil.copytree(Source_path_2, final)\n",
        "\n",
        "# saves final result in the folder named \"final result\"\n",
        "# I created it in two folders to make the run from the notebook easier for the user at the beginning to be able to do the quality controle with a detemined folder name \"final_result\"\n",
        "today = datetime.now()\n",
        "date_info = today.strftime('%Y%m%d_%H%M%S')\n",
        "final_date = Saving_path + f\"/{date_info}_final_result\"\n",
        "shutil.copytree(final, final_date)\n",
        "\n",
        "# remove unnecessary folders\n",
        "for path in remove_folder:\n",
        "  shutil.rmtree(path)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OPKfyCwrut2b",
        "cellView": "form"
      },
      "source": [
        "\n",
        "#@markdown Provide a folder path on your gdrive to save the DAIN-interpolation result\n",
        "Save_path = \"/content/sample_data\" #@param {type:\"string\"}\n",
        "\n",
        "if not os.path.exists(Save_path):\n",
        "  os.mkdir(Save_path)\n",
        "\n",
        "folder_name = final_date.split(\"/\")[-1]\n",
        "destination = os.path.join(Save_path, folder_name)\n",
        "shutil.copytree(final_date, destination)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tta4phe6VMtO"
      },
      "source": [
        "\n",
        "# **6. Error mapping and quality metrics estimation**\n",
        "---"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dbeEauCL7YjG"
      },
      "source": [
        "**WORKFLOW**\n",
        "\n",
        "- You can compare the quality of first downsampled images (performed in section 3) and re-interpolated with the ground truth image used for down-scaling\n",
        "- The first cells loads the images and shows the dimensions\n",
        "- The second cell evaluates the image and saves 2 csv files showing the mean quality values and the other one showing the quality values for each slice. It also saves the SSIM and RMSE maps of the dataset in the same folder\n",
        "- The third cell shows you SSIM and RMSE maps of a selected image slice"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4_eedldMK8bs"
      },
      "source": [
        "#### Explanation of evaluation metrics\n",
        "<font size = 3>This option can be used to compare the downsampled-z or downsample-t or zoom option with the original file to estimate the performance of the network\n",
        "---\n",
        "\n",
        "<font size = 2>This section will display SSIM maps and RSE maps as well as calculating total SSIM, NRMSE and PSNR metrics for all the images provided in the \"GT_Image_QC\" and \"Prediction_Image_QC\" !\n",
        "\n",
        "<font size = 3>**1. The SSIM (structural similarity) map** \n",
        "\n",
        "<font size = 2>The SSIM metric is used to evaluate whether two images contain the same structures. It is a normalized metric and an SSIM of 1 indicates a perfect similarity between two images. Therefore for SSIM, the closer to 1, the better. The SSIM maps are constructed by calculating the SSIM metric in each pixel by considering the surrounding structural similarity in the neighbourhood of that pixel (currently defined as window of 11 pixels and with Gaussian weighting of 1.5 pixel standard deviation, see our Wiki for more info). \n",
        "\n",
        "<font size=2>**mSSIM** is the SSIM value calculated across the entire window of both images.\n",
        "\n",
        "<font size=2>**The output below shows the SSIM maps with the mSSIM**\n",
        "\n",
        "<font size = 3>**2. The RSE (Root Squared Error) map** \n",
        "\n",
        "<font size = 2>This is a display of the root of the squared difference between the normalized predicted and target or the source and the target. In this case, a smaller RSE is better. A perfect agreement between target and prediction will lead to an RSE map showing zeros everywhere (dark).\n",
        "\n",
        "\n",
        "<font size =2>**NRMSE (normalised root mean squared error)** gives the average difference between all pixels in the images compared to each other. Good agreement yields low NRMSE scores.\n",
        "\n",
        "<font size = 2>**PSNR (Peak signal-to-noise ratio)** is a metric that gives the difference between the ground truth and prediction (or source input) in decibels, using the peak pixel values of the prediction and the MSE between the images. The higher the score the better the agreement.\n",
        "\n",
        "<font size=2>**The output below shows the RSE maps with the NRMSE and PSNR values.**\n",
        "\n",
        "<font size=2>**Quality check just works for not RGB images**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vWGw4s-RK8bt",
        "cellView": "form"
      },
      "source": [
        "#@title ####**6.1 Check for image shift**\n",
        "#@markdown Select on sample predicted image and one GT image for comparison. This cell compares \n",
        "### Check for a channel shift of the reconstructed image\n",
        "import sys\n",
        "sys.path.insert(0,'/content/DAIN/load_functions')\n",
        "from quality_metrics_estimation import correct_channels\n",
        "from quality_metrics_estimation import bcolors\n",
        "from quality_metrics_estimation import create_shift_image\n",
        "from quality_metrics_estimation import norm_minmse\n",
        "from quality_metrics_estimation import structural_similarity\n",
        "from quality_metrics_estimation import psnr\n",
        "from quality_metrics_estimation import get_full_file_paths\n",
        "from skimage import io\n",
        "import pandas as pd\n",
        "import numpy as np\n",
        "# program a channel_shift\n",
        "Prediction_Image_QC = \"/content/DAIN/demo/final_result\" #@param{type:\"string\"}\n",
        "GT_Image_QC = \"/content/DAIN/demo/sim_particles_2DF_512\" #@param{type:\"string\"}\n",
        "\n",
        "# read out all images in the selected folder\n",
        "prediction_flist_paths = get_full_file_paths(Prediction_Image_QC)\n",
        "GT_flist_paths = get_full_file_paths(GT_Image_QC)\n",
        "prediction_flist_paths.sort()\n",
        "GT_flist_paths.sort()\n",
        "\n",
        "Predicted_image = prediction_flist_paths[-1]\n",
        "GT_Image = GT_flist_paths[-1]\n",
        "\n",
        "#@markdown You can either select a given image based on the t and z-slice. If unticked the center image of the z and z dimension will be selected automatically.\n",
        "select_slices = True #@param {type:\"boolean\"}\n",
        "\n",
        "int_img = io.imread(Predicted_image)\n",
        "int_img, use_RGB = correct_channels(int_img)\n",
        "print(int_img.shape)\n",
        "\n",
        "if select_slices:\n",
        "  t_slice = 7 #@param {type:\"raw\"}\n",
        "  z_slice = 0 #@param {type:\"raw\"}\n",
        "  channel = 0 #@param {type:\"raw\"}\n",
        "else:\n",
        "  t_slice = t//2\n",
        "  z_slice = z//2\n",
        "  channel = 0\n",
        "\n",
        "if use_RGB:\n",
        "  t, z, y_dim, x_dim, c= int_img.shape\n",
        "  int_img = int_img[t_slice, z_slice,:,:,channel]\n",
        "\n",
        "else:\n",
        "  t, z, y_dim, x_dim = int_img.shape\n",
        "  int_img = int_img[t_slice, z_slice,:,:]\n",
        "\n",
        "gt_img = io.imread(GT_Image)\n",
        "gt_img, use_RGB = correct_channels(gt_img)\n",
        "\n",
        "print(gt_img.shape)\n",
        "if use_RGB:\n",
        "  t, z, y_dim, x_dim, c= gt_img.shape\n",
        "  gt_img = gt_img[t_slice, z_slice,:,:,channel]\n",
        "\n",
        "else:\n",
        "  t, z, y_dim, x_dim = gt_img.shape\n",
        "  gt_img = gt_img[t_slice, z_slice,:,:]\n",
        "\n",
        "# try different shifts\n",
        "x_list = [-4,-3,-2,-1,0,1,2,3,4]\n",
        "y_list = [-4,-3,-2,-1,0,1,2,3,4]\n",
        "dict_NRMSE_values = {}\n",
        "dict_SSIM_values = {}\n",
        "dict_PSNR_values = {}\n",
        "\n",
        "for y_shift in y_list:\n",
        "  for x_shift in x_list:\n",
        "    # print(temp_img.shape)\n",
        "    shifted_img = create_shift_image(int_img, y_shift, x_shift)\n",
        "        \n",
        "    # evaluate SSIM, PSNR, NRMSE\n",
        "    test_GT_norm, test_prediction_norm = norm_minmse(gt_img, shifted_img, normalize_gt=True)\n",
        "    img_RSE_GTvsPrediction = np.sqrt(np.square(test_GT_norm - test_prediction_norm))\n",
        "    NRMSE_GTvsPrediction = np.sqrt(np.mean(img_RSE_GTvsPrediction))\n",
        "    index_SSIM_GTvsPrediction, img_SSIM_GTvsPrediction = structural_similarity(test_GT_norm, test_prediction_norm, data_range=1.0, full=True)\n",
        "    PSNR_GTvsPrediction = psnr(test_GT_norm, test_prediction_norm, data_range=1.0)\n",
        "\n",
        "\n",
        "    dict_NRMSE_values[str(y_shift)+ \"_\" +str(x_shift)] = [NRMSE_GTvsPrediction]\n",
        "    dict_SSIM_values[str(y_shift)+ \"_\" +str(x_shift)] = [index_SSIM_GTvsPrediction]\n",
        "    dict_PSNR_values[str(y_shift)+ \"_\" +str(x_shift)] = [PSNR_GTvsPrediction]\n",
        "\n",
        "    \n",
        "    # print(f\"y_shift {y_shift}, x_shift  {x_shift}\")\n",
        "    # print(f\"NRMSE_GTvsPrediction {NRMSE_GTvsPrediction}, index_SSIM_GTvsPrediction {index_SSIM_GTvsPrediction}, PSNR_GTvsPrediction {PSNR_GTvsPrediction}\")\n",
        "\n",
        "## uncommented - COMMENT: removed it because confusing for the user. just useful if for some reason the image gets shifted for some pixels in the x or y direction. then this could be uncommented to find the number of shifted pixels and correct it accordingly.\n",
        "# import pylab as pl\n",
        "# # pixel_shift evaluation\n",
        "# pd.DataFrame.from_dict(dict_SSIM_values).transpose().plot.bar(figsize=(25,2))\n",
        "# print(f\"best conditions according to SSIM: {max(dict_SSIM_values, key=dict_SSIM_values.get)}\")\n",
        "# pl.suptitle(\"Different values for SSIM accoring to different x-y shifts - max best\")\n",
        "\n",
        "\n",
        "# pd.DataFrame.from_dict(dict_PSNR_values).transpose().plot.bar(figsize=(25,2))\n",
        "# print(f\"best conditions according to PSNR: {max(dict_PSNR_values, key=dict_PSNR_values.get)}\")\n",
        "# pl.suptitle(\"Different values for PSNR accoring to different x-y shifts - max best\")\n",
        "\n",
        "# pd.DataFrame.from_dict(dict_NRMSE_values).transpose().plot.bar(figsize=(25,2))\n",
        "# print(f\"best conditions according to NRMSE: {min(dict_NRMSE_values, key=dict_NRMSE_values.get)}\")\n",
        "# pl.suptitle(\"Different values for PSNR accoring to different y-x shifts - min best\")\n",
        "\n",
        "\n",
        "\n",
        "#################################################################\n",
        "#@title ####**4.3.3 Perform quality control over all files**\n",
        "#@markdown Choose the folders to compare original with predicted images with the ground truth images for quality control\n",
        "from quality_metrics_estimation import  make_folder_with_date\n",
        "from quality_metrics_estimation import get_full_file_paths\n",
        "import csv\n",
        "import os\n",
        "\n",
        "# Prediction_Image_QC = \"/content/gdrive/MyDrive/1.2_BIG_DATA_PhD_Project_2/0.PAPER_EXPERIMENT/20210321_Tchern_dataset_Usecase/Zoomed_512_T/128_13_BIC_2x\" #@param{type:\"string\"}\n",
        "# GT_Image_QC = \"/content/gdrive/MyDrive/1.2_BIG_DATA_PhD_Project_2/0.PAPER_EXPERIMENT/20210321_Tchern_dataset_Usecase/Zoomed_512_T/256_13\" #@param{type:\"string\"}\n",
        "\n",
        "# Create quality control folder\n",
        "base_folder = \"/\".join(Prediction_Image_QC.split(\"/\")[:-1])\n",
        "exp_name = Prediction_Image_QC.split(\"/\")[-1]\n",
        "quality_test_folder = os.path.join(base_folder, f\"quality_tests-{exp_name}\")\n",
        "\n",
        "print(quality_test_folder)\n",
        "if os.path.exists(quality_test_folder)==False:\n",
        "  os.mkdir(quality_test_folder)\n",
        "# quality_test_folder = make_folder_with_date(base_folder, f\"{exp_name}_quality_tests\")\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "if len(prediction_flist_paths) != len(GT_flist_paths):\n",
        "  print(f\"{bcolors.WARNING}There is an unequal number of images in the two selected folders{bcolors.ENDC}\")\n",
        "else:\n",
        "  print(f\"There are {len(GT_flist_paths)} files in both folders\")\n",
        "\n",
        "file_num = len(GT_flist_paths)\n",
        "\n",
        "for num in range(file_num):\n",
        "    # create log-csv files for the mean of each stack and the similarity/error within each image slice\n",
        "    with open(quality_test_folder +f'/{exp_name}_every_slice{num}.csv', 'w', newline='') as file:\n",
        "          writer = csv.writer(file)\n",
        "          writer.writerow([\"File\", \"Z-Dimension\", \"T-Dimension\", \"index_SSIM_GTvsPrediction\", \"NRMSE_GTvsPrediction\",\"PSNR_GTvsPrediction\"])\n",
        "\n",
        "    with open(quality_test_folder +f'/{exp_name}_mean{num}.csv', 'w', newline='') as file:\n",
        "          writer = csv.writer(file)\n",
        "          writer.writerow([\"File\", \"mSSIM_GvP_mean\", \"NRMSE_GvP_mean\", \"PSNR_GvP_mean\"])\n",
        "\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MUKo-Bna4fFF",
        "cellView": "form"
      },
      "source": [
        "#@title ####**6.2 Run quality control**\n",
        "\n",
        "#@markdown This cells saves the evaluation images and values for SSIM, PSNR, MSR in \"*.csv\" files.\n",
        "\n",
        "#@markdown You can either select a manual shift for the x and y-dimension of the images. If unticked the best conditions ecaluated in **Section 3.4.2** will be take.\n",
        "\n",
        "from tqdm import tqdm\n",
        "import os\n",
        "\n",
        "select_manual_shift = True\n",
        "y_shift = 0\n",
        "x_shift = 0\n",
        "#@markdown If multi channel image choose the channel of interest:\n",
        "channel = 0 #@param {type:\"raw\"}\n",
        "\n",
        "## uncommented - COMMENT: removed it because confusing for the user. just useful if for some reason the image gets shifted for some pixels in the x or y direction. then this could be uncommented to find the number of shifted pixels and correct it accordingly.\n",
        "\n",
        "# if select_manual_shift:\n",
        "#   y_shift = 0 #@param {type:\"raw\"}\n",
        "#   x_shift = 0 #@param {type:\"raw\"}\n",
        "#   channel = 0 #@param {type:\"raw\"}\n",
        "\n",
        "# else:\n",
        "#   y_shift = int(max(dict_PSNR_values, key=dict_PSNR_values.get).split(\"_\")[0])\n",
        "#   x_shift = int(max(dict_PSNR_values, key=dict_PSNR_values.get).split(\"_\")[1])\n",
        "#   channel = 0 \n",
        "\n",
        "\n",
        "# These lists will be used to collect all the metrics values per slice\n",
        "file_name_list = []\n",
        "T_slice_number_list = []\n",
        "Z_slice_number_list = []\n",
        "mSSIM_GvP_list = []\n",
        "NRMSE_GvP_list = []\n",
        "PSNR_GvP_list = []\n",
        "\n",
        "\n",
        "img_GT = io.imread(GT_flist_paths[0])\n",
        "img_GT, use_RGB = correct_channels(img_GT)\n",
        "if use_RGB:\n",
        "  img_GT = img_GT[:,:,:,:,channel]\n",
        "\n",
        "\n",
        "img_prediction = io.imread(prediction_flist_paths[0])\n",
        "img_prediction, use_RGB = correct_channels(img_prediction)\n",
        "if use_RGB:\n",
        "  img_prediction = img_prediction[:,:,:,:,channel]\n",
        "\n",
        "# img_prediction = img_prediction.get_image_data(\"TZXY\")\n",
        "\n",
        "t_slices = img_prediction.shape[0]\n",
        "z_slices = img_prediction.shape[1]\n",
        "print(img_GT.shape)\n",
        "print(img_prediction.shape)\n",
        "print(f\"The t-dimension is {t_slices}, and the z-dimension is {z_slices}\")\n",
        "\n",
        "#  Create a zero matrix as template for the images\n",
        "img_SSIM_GTvsPrediction_stack = np.zeros((t_slices, z_slices, img_prediction.shape[2], img_prediction.shape[3]))\n",
        "img_RSE_GTvsPrediction_stack = np.zeros((t_slices, z_slices, img_prediction.shape[2], img_prediction.shape[3]))\n",
        "\n",
        "# compare each image - slice per slice and save the calculation in csvs and as new image\n",
        "\n",
        "num = 0 \n",
        "for pred_file_path, GT_file_path in zip(prediction_flist_paths, GT_flist_paths):\n",
        "    # load the two images\n",
        "    img_prediction = io.imread(pred_file_path)\n",
        "    img_prediction, use_RGB = correct_channels(img_prediction)\n",
        "    if use_RGB:\n",
        "      img_prediction = img_prediction[:,:,:,:,channel]\n",
        "\n",
        "    # img_prediction = img_prediction.get_image_data(\"TZXY\")\n",
        "    img_GT = io.imread(GT_file_path)\n",
        "    img_GT, use_RGB = correct_channels(img_GT)\n",
        "    if use_RGB:\n",
        "      img_GT = img_GT[:,:,:,:,channel]\n",
        "\n",
        "    # img_GT = img_GT.get_image_data(\"TZXY\")\n",
        "    file_name = \"/\".join(GT_file_path.split(\"/\")[-1:])\n",
        "    for z in tqdm(range(z_slices)): \n",
        "        for t in range(t_slices): \n",
        "\n",
        "          # --------------------------- perform optimal shift---------------------------\n",
        "          \n",
        "          one_img_prediction = img_prediction[t][z]\n",
        "          one_img_prediction = create_shift_image(one_img_prediction, y_shift, x_shift)\n",
        "\n",
        "          # -------------------------------- Prediction --------------------------------\n",
        "\n",
        "          test_GT_norm, test_prediction_norm = norm_minmse(img_GT[t][z], one_img_prediction, normalize_gt=True)\n",
        "\n",
        "          # ------------------------ Calculate the SSIM metric and maps ----------------\n",
        "          # Calculate the SSIM maps and index\n",
        "          index_SSIM_GTvsPrediction, img_SSIM_GTvsPrediction = structural_similarity(test_GT_norm, test_prediction_norm, data_range=1.0, full=True)\n",
        "\n",
        "          #Calculate ssim_maps\n",
        "          # img_SSIM_GTvsPrediction_stack[t][z] = np.float32(img_SSIM_GTvsPrediction)\n",
        "          img_SSIM_GTvsPrediction_stack[t][z] = img_SSIM_GTvsPrediction\n",
        "\n",
        "\n",
        "          # ----------------------- Calculate the NRMSE metrics ------------------------\n",
        "          # Calculate the Root Squared Error (RSE) maps\n",
        "          img_RSE_GTvsPrediction = np.sqrt(np.square(test_GT_norm - test_prediction_norm))\n",
        "\n",
        "          # Calculate SE maps\n",
        "          img_RSE_GTvsPrediction_stack[t][z] = np.float32(img_RSE_GTvsPrediction)\n",
        "\n",
        "\n",
        "          # Normalised Root Mean Squared Error (here it's valid to take the mean of the image)\n",
        "          NRMSE_GTvsPrediction = np.sqrt(np.mean(img_RSE_GTvsPrediction))\n",
        "\n",
        "          # Calculate the PSNR between the images\n",
        "          PSNR_GTvsPrediction = psnr(test_GT_norm, test_prediction_norm, data_range=1.0)\n",
        "\n",
        "          with open(quality_test_folder +f'/{exp_name}_every_slice{num}.csv', 'a', newline='') as file:\n",
        "              writer = csv.writer(file)\n",
        "              writer.writerow([file_name, z, t, index_SSIM_GTvsPrediction, NRMSE_GTvsPrediction,PSNR_GTvsPrediction])\n",
        "\n",
        "          mSSIM_GvP_list.append(index_SSIM_GTvsPrediction)\n",
        "          NRMSE_GvP_list.append(NRMSE_GTvsPrediction)\n",
        "          PSNR_GvP_list.append(PSNR_GTvsPrediction)\n",
        "\n",
        "    # ----------- Change the stacks to 32 bit images and fix channels -----------\n",
        "    from skimage import img_as_float32\n",
        "    img_SSIM_GTvsPrediction_stack_32 = img_as_float32(img_SSIM_GTvsPrediction_stack, force_copy=False)\n",
        "    # img_SSIM_GTvsPrediction_stack_32= reshape_data(img_SSIM_GTvsPrediction_stack_32, \"TZXY\",\"STCZXY\")\n",
        "\n",
        "    img_RSE_GTvsPrediction_stack_32 = img_as_float32(img_RSE_GTvsPrediction_stack, force_copy=False)\n",
        "    # img_RSE_GTvsPrediction_stack_32= reshape_data(img_RSE_GTvsPrediction_stack_32, \"TZXY\",\"STCZXY\")\n",
        "\n",
        "\n",
        "    # ----------- Saving the error map stacks -----------\n",
        "    path_file = os.path.join(quality_test_folder, file_name)\n",
        "    SSIM_path = path_file[:-4] + \"SSIM_GTvsPrediction.tif\"\n",
        "    RSE_path =  path_file[:-4] + 'RSE_GTvsPrediction.tif'\n",
        "\n",
        "    io.imsave(SSIM_path,img_SSIM_GTvsPrediction_stack_32)\n",
        "    io.imsave(RSE_path,img_RSE_GTvsPrediction_stack_32)\n",
        "\n",
        "    # ----------- Saving mean errors in CSV -----------\n",
        "    mSSIM_GvP_mean = (sum(mSSIM_GvP_list)/len(mSSIM_GvP_list))\n",
        "    NRMSE_GvP_mean = (sum(NRMSE_GvP_list)/len(NRMSE_GvP_list))\n",
        "    PSNR_GvP_mean = (sum(PSNR_GvP_list)/len(PSNR_GvP_list))\n",
        "    with open(quality_test_folder + f'/{exp_name}_mean{num}.csv', 'a', newline='') as file:\n",
        "        writer = csv.writer(file)\n",
        "        writer.writerow([file_name, mSSIM_GvP_mean, NRMSE_GvP_mean, PSNR_GvP_mean])\n",
        "    num += 1\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FK2iuVKl4iXV",
        "cellView": "form"
      },
      "source": [
        "#@title ####**6.3 Visualize single images from a selected file**\n",
        "\n",
        "#@markdown Select a file number (based on the number of image that were processed before) and the t/z slices of images to compare\n",
        "#@markdown (Index starts with 0)\n",
        "import matplotlib.pyplot as plt\n",
        "from quality_metrics_estimation import bcolors\n",
        "\n",
        "# nr_file =  0#@param {type:\"integer\"}\n",
        "z_dimension =  0#@param {type:\"integer\"}\n",
        "t_dimension =  3#@param {type:\"integer\"}\n",
        "\n",
        "select_manual_shift = True\n",
        "y_shift = 0\n",
        "x_shift = 0\n",
        "#@markdown If multi channel image choose the channel of interest:\n",
        "channel = 0 #@param {type:\"raw\"}\n",
        "\n",
        "## uncommented - COMMENT: removed it because confusing for the user. just useful if for some reason the image gets shifted for some pixels in the x or y direction. then this could be uncommented to find the number of shifted pixels and correct it accordingly.\n",
        "## @markdown You can either select a manual shift for the x and y-dimension of the images. If unticked the best conditions ecaluated in **Section 3.4.2** will be take.\n",
        "\n",
        "# select_manual_shift = True #@param {type:\"boolean\"}\n",
        "# if select_manual_shift:\n",
        "#   y_shift = 0 #@param {type:\"raw\"}\n",
        "#   x_shift = 0 #@param {type:\"raw\"}\n",
        "#   # channel = 0 #@param {type:\"integer\"}\n",
        "\n",
        "# else:\n",
        "#   y_shift = int(max(dict_PSNR_values, key=dict_PSNR_values.get).split(\"_\")[0])\n",
        "#   x_shift = int(max(dict_PSNR_values, key=dict_PSNR_values.get).split(\"_\")[1])\n",
        "\n",
        "display_all_files =True\n",
        "if display_all_files== True:\n",
        "   for nr_file in range(file_num):\n",
        "        Prediction_Image_QC = prediction_flist_paths[nr_file]\n",
        "        GT_Image_QC = GT_flist_paths[nr_file]\n",
        "        # get file names of both\n",
        "        file_name_pred = \"/\".join(Prediction_Image_QC.split(\"/\")[-1:])\n",
        "        file_name_GT = \"/\".join(GT_Image_QC.split(\"/\")[-1:])\n",
        "        # load the two images\n",
        "        img_prediction = io.imread(Prediction_Image_QC)\n",
        "        img_prediction, use_RGB = correct_channels(img_prediction)\n",
        "        if use_RGB:\n",
        "          img_prediction = img_prediction[:,:,:,:,channel]\n",
        "\n",
        "        # img_prediction = img_prediction.get_image_data(\"TZXY\")\n",
        "        prediction_dim =img_prediction.shape\n",
        "        print(f\"Source dimension: {prediction_dim}\")\n",
        "\n",
        "        img_GT = io.imread(GT_Image_QC)\n",
        "        img_GT, use_RGB = correct_channels(img_GT)\n",
        "        if use_RGB:\n",
        "          img_GT = img_GT[:,:,:,:,channel]\n",
        "\n",
        "        # img_GT = img_GT.get_image_data(\"TZXY\")\n",
        "        GT_dim =img_GT.shape\n",
        "        print(f\"Target dimension: {GT_dim}\")\n",
        "\n",
        "        #prepare slice\n",
        "        prediction_img_slice = img_prediction[t_dimension,z_dimension,:,:]\n",
        "        GT_img_slice = img_GT[t_dimension,z_dimension,:,:]\n",
        "        print(GT_img_slice.shape)\n",
        "\n",
        "\n",
        "        # -------------------------------- Calculate the SSIM metric and maps --------------------------------\n",
        "\n",
        "        img_GT_slice = img_GT[t_dimension][z_dimension]\n",
        "        img_prediction_slice = img_prediction[t_dimension][z_dimension]\n",
        "        img_prediction_slice = create_shift_image(img_prediction_slice, y_shift, x_shift)\n",
        "\n",
        "        test_GT_norm, test_prediction_norm = norm_minmse(img_GT_slice, img_prediction_slice, normalize_gt=True)\n",
        "\n",
        "        # Calculate the SSIM maps and index data_range=test_prediction_norm.max() - test_prediction_norm.min()\n",
        "        index_SSIM_GTvsPrediction, img_SSIM_GTvsPrediction = structural_similarity(test_GT_norm, test_prediction_norm, data_range=1, full=True)\n",
        "        img_SSIM_GTvsPrediction = np.float32(img_SSIM_GTvsPrediction)\n",
        "\n",
        "        # Calculate the RSE\n",
        "        img_RSE_GTvsPrediction = np.sqrt(np.square(test_GT_norm - test_prediction_norm))\n",
        "        img_RSE_GTvsPrediction = np.float32(img_RSE_GTvsPrediction)\n",
        "\n",
        "        # Normalised Root Mean Squared Error (here it's valid to take the mean of the image)\n",
        "        NRMSE_GTvsPrediction = np.sqrt(np.mean(img_RSE_GTvsPrediction))\n",
        "\n",
        "        # Calculate the PSNR between the images\n",
        "        PSNR_GTvsPrediction = psnr(test_GT_norm, test_prediction_norm, data_range=1.0)\n",
        "\n",
        "\n",
        "        # -------------------------------- Plotting --------------------------------\n",
        "\n",
        "        plt.figure(figsize=(15,15))\n",
        "        # Currently only displays the last computed set, from memory\n",
        "\n",
        "        #Setting up colours\n",
        "        cmap = plt.cm.Greys\n",
        "        cmap = cmap.reversed()\n",
        "\n",
        "        # Target (Ground-truth)\n",
        "        plt.subplot(2,3,1)\n",
        "        plt.tick_params(\n",
        "            axis='both',      # changes apply to the x-axis and y-axis\n",
        "            which='both',      # both major and minor ticks are affected\n",
        "            bottom=False,      # ticks along the bottom edge are off\n",
        "            top=False,        # ticks along the top edge are off\n",
        "            left=False,       # ticks along the left edge are off\n",
        "            right=False,         # ticks along the right edge are off\n",
        "            labelbottom=False,\n",
        "            labelleft=False)  \n",
        "        plt.xlabel(f'image {file_name}')\n",
        "        plt.imshow(img_GT_slice, cmap = cmap)\n",
        "        plt.title('GT (slice #'+str(t_dimension)+'-'+str(t_dimension)+')')\n",
        "        \n",
        "\n",
        "        # Source\n",
        "        plt.subplot(2,3,2)\n",
        "        plt.tick_params(\n",
        "            axis='both',      # changes apply to the x-axis and y-axis\n",
        "            which='both',      # both major and minor ticks are affected\n",
        "            bottom=False,      # ticks along the bottom edge are off\n",
        "            top=False,        # ticks along the top edge are off\n",
        "            left=False,       # ticks along the left edge are off\n",
        "            right=False,         # ticks along the right edge are off\n",
        "            labelbottom=False,\n",
        "            labelleft=False)  \n",
        "        plt.xlabel(f'image {file_name}')\n",
        "        plt.imshow(img_prediction_slice, aspect='equal', cmap = cmap)\n",
        "        plt.title('Prediction (slice #'+str(t_dimension)+'-'+str(t_dimension)+')')\n",
        "\n",
        "\n",
        "        #Setting up colours\n",
        "        plt.figure(figsize=(15,15))\n",
        "        cmap = plt.cm.CMRmap\n",
        "\n",
        "        # img_SSIM_GTvsPrediction\n",
        "        plt.subplot(2,3,1)\n",
        "        plt.tick_params(\n",
        "            axis='both',      # changes apply to the x-axis and y-axis\n",
        "            which='both',      # both major and minor ticks are affected\n",
        "            bottom=False,      # ticks along the bottom edge are off\n",
        "            top=False,        # ticks along the top edge are off\n",
        "            left=False,       # ticks along the left edge are off\n",
        "            right=False,         # ticks along the right edge are off\n",
        "            labelbottom=False,\n",
        "            labelleft=False)  \n",
        "        plt.xlabel(f'image {file_name}')\n",
        "        img_SSIM_GTvsPrediction_plt = plt.imshow(img_SSIM_GTvsPrediction, cmap = cmap, vmin=0,vmax=1)\n",
        "        plt.colorbar(img_SSIM_GTvsPrediction_plt,fraction=0.046, pad=0.04)\n",
        "        plt.title('iSSIM map: GT vs. Prediction (slice #'+str(t_dimension)+'-'+str(t_dimension)+')')\n",
        "\n",
        "        # img_RSE_GTvsPrediction\n",
        "        plt.subplot(2,3,2)\n",
        "        plt.tick_params(\n",
        "            axis='both',      # changes apply to the x-axis and y-axis\n",
        "            which='both',      # both major and minor ticks are affected\n",
        "            bottom=False,      # ticks along the bottom edge are off\n",
        "            top=False,        # ticks along the top edge are off\n",
        "            left=False,       # ticks along the left edge are off\n",
        "            right=False,         # ticks along the right edge are off\n",
        "            labelbottom=False,\n",
        "            labelleft=False)  \n",
        "        plt.xlabel(f'image {file_name}')\n",
        "        RSE_GTvsPrediction_plt = plt.imshow(img_RSE_GTvsPrediction, aspect='equal', cmap = cmap, vmin=0,vmax=1)\n",
        "        plt.colorbar(RSE_GTvsPrediction_plt,fraction=0.046, pad=0.04)\n",
        "        plt.title('img_RSE_GTvsPrediction (slice #'+str(t_dimension)+'-'+str(t_dimension)+')')\n",
        "\n",
        "        print(f\"{bcolors.WARNING}PSNR_GTvsPrediction:       {PSNR_GTvsPrediction}{bcolors.ENDC}\")\n",
        "        print(f\"{bcolors.WARNING}NRMSE_GTvsPrediction:      {NRMSE_GTvsPrediction}{bcolors.OKBLUE}\")\n",
        "        print(f\"{bcolors.WARNING}index_SSIM_GTvsPrediction: {index_SSIM_GTvsPrediction}{bcolors.OKGREEN}\")\n",
        "\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-fml8tRkrLyq"
      },
      "source": [
        "#Good luck!"
      ]
    }
  ]
}