exercise.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1699d1c2",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "# A Generative Modelling Approach to Image translation\n",
    "Written by [Samuel Tonks](https://github.com/Tonks684), Krull Lab University of Birmingham UK, with many inputs and bugfixes from [Eduardo Hirata-Miyasaki](https://github.com/edyoshikun), [Ziwen Liu](https://github.com/ziw-liu) and [Shalin Mehta](https://github.com/mattersoflight) of CZ Biohub San Francisco."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0ea5d1b2",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 0
   },
   "source": [
    "## Overview\n",
    "\n",
    "In part 2 of the image_translation exercise, we will predict fluorescence images of nuclei markers only from quantitative phase images of cells, using a specific type of generative model called a Conditional Generative Adversarial Network (cGAN). In contrast to a regression-based approach, cGANs learn to map from the phase contrast domain to a distirbution of virtual staining solutions. In this work we will utilise the [Pix2PixHD GAN](https://arxiv.org/abs/1711.11585) used in our recent [virtual staining works](https://ieeexplore.ieee.org/abstract/document/10230501?casa_token=NEyrUDqvFfIAAAAA:tklGisf9BEKWVjoZ6pgryKvLbF6JyurOu5Jrgoia1QQLpAMdCSlP9gMa02f3w37PvVjdiWCvFhA). For more details on the architecture and loss components of cGANs and Pix2PixHD GAN please see the READ.me. \n",
    "\n",
    "During this exercise will assess the different loss components of a pre-trained Pix2PixHD for the virtual nuclei staining task. We will then evaluate the performance of the model on unseen data using the same pixel-level and instance-level metrics as in the previous section. We will compare the performance of the Pix2PixHD GAN with the regression-based model Viscy. Finally, as a bonus, we will explore the variability and uncertainty in the virtual stain predictions using [MC-Dropout](https://arxiv.org/abs/1506.02142).\n",
    "\n",
    "## References\n",
    "- [Wang, T. et al. (2018) High-resolution image synthesis and semantic manipulation with conditional GANs](https://arxiv.org/abs/1711.11585)\n",
    "- [Tonks, S. et al. (2023) Evaluating virtual staining for high-throughput screening](https://ieeexplore.ieee.org/abstract/document/10230501?casa_token=NEyrUDqvFfIAAAAA:tklGisf9BEKWVjoZ6pgryKvLbF6JyurOu5Jrgoia1QQLpAMdCSlP9gMa02f3w37PvVjdiWCvFhA)\n",
    "- [Gal, Y. et al. (2016) Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning](https://arxiv.org/abs/1506.02142)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54094c01",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 2
   },
   "source": [
    "## Goals\n",
    "\n",
    "This part of the exercise is organized in 5 parts.<br>\n",
    "\n",
    "As you have already explored the data in the previous parts, we will focus on training and evaluating Pix2PixHD GAN. The parts are as follows:<br>\n",
    "\n",
    "* **Part 1** - Define dataloaders & walk through steps to train a Pix2PixHD GAN.<br>\n",
    "* **Part 2** - Load and assess a pre-trained Pix2PixGAN using tensorboard, discuss the different loss components and how new hyper-parameter configurations could impact performance.<br>\n",
    "* **Part 3** - Evaluate performance of pre-trained Pix2PixGAN using pixel-level and instance-level metrics.<br>\n",
    "* **Part 4** - Compare the performance of Viscy (regression-based) with Pix2PixHD GAN (generative modelling approach)<br>\n",
    "* **Part 5** - *BONUS*: Sample different virtual staining solutions from the Pix2PixHD GAN using [MC-Dropout](https://arxiv.org/abs/1506.02142) and explore the variability and subsequent uncertainty in the virtual stain predictions.<br>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce7825e3",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 0
   },
   "source": [
    "<div class=\"alert alert-danger\">\n",
    "Set your python kernel to <span style=\"color:black;\">06_image_translation</span>\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e66da83b",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 0
   },
   "source": [
    "If you have issues with getting the kernel to load please follow the following steps:\n",
    "1. Ctrl+Shift+P to open Command Palette. Type Python: Select Interpreter and select 06_image_translation.\n",
    "2. Register the environment as a Kernel using the below line of code.\n",
    "3. Reload VS Code via Ctrl+Shift+P, then select Reload Window."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0cc50a47",
   "metadata": {},
   "outputs": [],
   "source": [
    "!python -m ipykernel install --user --name 06_image_translation --display-name \"Python 06_image_translation\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a4a7e68",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 0
   },
   "source": [
    "# Part 1: Define dataloaders & walk through steps to train a Pix2PixHD GAN.\n",
    "---------\n",
    "The focus of this part of the exercise is on understanding a generative modelling approach to image translation, how to train and evaluate a cGAN, and explore some hyperparameters of the cGAN. \n",
    "\n",
    "Learning goals:\n",
    "\n",
    "- Load dataset and configure dataloader.\n",
    "- Configure Pix2PixHD GAN to train for translating from phase to nuclei.\n",
    "\n",
    "\n",
    "Before we start please set the first section of the parent_dir to your personal path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3712f6b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TO DO: Change the path to the directory where the data and code is stored is stored.\n",
    "import os\n",
    "import sys\n",
    "parent_dir = os.path.abspath(\"ADD_HERE/data/06_image_translation/part2/GAN_code/GANs_MI2I/\")\n",
    "sys.path.append(parent_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50c9db51",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "import torch\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from tqdm import tqdm\n",
    "from skimage import metrics\n",
    "from tifffile import imread, imsave\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.animation as animation\n",
    "from IPython.display import HTML\n",
    "\n",
    "from cellpose import models\n",
    "from typing import List, Tuple\n",
    "from numpy.typing import ArrayLike\n",
    "\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "\n",
    "# Import all the necessary hyperparameters and configurations for training.\n",
    "from pix2pixHD.options.train_options import TrainOptions\n",
    "from pix2pixHD.options.test_options import TestOptions\n",
    "\n",
    "# Import Pytorch dataloader and transforms.\n",
    "from pix2pixHD.data.data_loader_dlmbl import CreateDataLoader\n",
    "\n",
    "# Import the model architecture.\n",
    "from pix2pixHD.models import create_model\n",
    "\n",
    "# Import helper functions for visualization and processing.\n",
    "from pix2pixHD.util.visualizer import Visualizer\n",
    "from pix2pixHD.util import util\n",
    "\n",
    "# Import train script.\n",
    "from pix2pixHD.train_dlmbl import train as train_model\n",
    "from pix2pixHD.test_dlmbl import inference as inference_model\n",
    "from pix2pixHD.test_dlmbl import sampling\n",
    "\n",
    "# pytorch lightning wrapper for Tensorboard.\n",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "\n",
    "# Import the same evaluation metrics as in the previous section.\n",
    "from viscy.evaluation.evaluation_metrics import mean_average_precision\n",
    "from torchmetrics.functional import accuracy, dice, jaccard_index\n",
    "\n",
    "# Initialize the default options and parse the arguments.\n",
    "opt = TrainOptions().parse()\n",
    "# Set the seed for reproducibility.\n",
    "util.set_seed(42)\n",
    "# Set the experiment folder name.\n",
    "translation_task = \"nuclei\"  # or \"cyto\" depending on your choice of target for virtual stain.\n",
    "opt.name = \"dlmbl_vsnuclei\"\n",
    "# Path to store all the logs.\n",
    "top_dir = Path(f\"~/data/06_image_translation/part2/\").expanduser() # TODO: Change this to point to your data directory.\n",
    "opt.checkpoints_dir = top_dir/\"GAN_code/GANs_MI2I/new_training_runs/\"\n",
    "Path(f'{opt.checkpoints_dir}/{opt.name}').mkdir(parents=True, exist_ok=True)\n",
    "output_image_folder = top_dir/\"tiff_files/\"\n",
    "# Initalize the tensorboard writer.\n",
    "writer = SummaryWriter(log_dir=opt.checkpoints_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "195de959",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 0
   },
   "source": [
    "## 1.1 Load Dataset & Configure Dataloaders.<br>\n",
    "Having already downloaded and split our training, validation and test sets we now need to load the data into the model. We will use the Pytorch DataLoader class to load the data in batches. The DataLoader class is an iterator that provides a consistent way to load data in batches. We will also use the CreateDataLoader class to load the data in the correct format for the Pix2PixHD GAN."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13b32758",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [],
   "source": [
    "# Initialize the Dataset and Dataloaders.\n",
    "\n",
    "## Define Dataset & Dataloader options.\n",
    "opt.dataroot = output_image_folder\n",
    "opt.data_type = 16  # Data type of the images.\n",
    "opt.loadSize = 512  # Size of the loaded phase image.\n",
    "opt.input_nc = 1  # Number of input channels.\n",
    "opt.output_nc = 1  # Number of output channels.\n",
    "opt.resize_or_crop = \"none\"  # Scaling and cropping of images at load time [resize_and_crop|crop|scale_width|scale_width_and_crop|none].\n",
    "opt.target = \"nuclei\"  # or \"cyto\" depending on your choice of target for virtual stain.\n",
    "\n",
    "# Load Training Set for input into model\n",
    "opt.isTrain = True \n",
    "train_dataloader = CreateDataLoader(opt)\n",
    "dataset_train = train_dataloader.load_data()\n",
    "print(f\"Total Training Images = {len(train_dataloader)}\")\n",
    "\n",
    "# Load Val Set\n",
    "opt.phase = \"val\"\n",
    "val_dataloader = CreateDataLoader(opt)\n",
    "dataset_val = val_dataloader.load_data()\n",
    "print(f\"Total Validation Images = {len(val_dataloader)}\")\n",
    "opt.phase= \"train\"\n",
    "\n",
    "# Plot a sample image from the training set."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10d33452",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 0
   },
   "source": [
    "## Configure Pix2PixHD GAN and train to predict nuclei from phase.\n",
    "Having loaded the data into the model we can now train the Pix2PixHD GAN to predict nuclei from phase. We will use the following hyperparameters to train the model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b5a7196b",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [],
   "source": [
    "# Define the parameters for the Generator.\n",
    "opt.ngf = 64  # Number of filters in the generator.\n",
    "opt.n_downsample_global = 4  # Number of downsampling layers in the generator.\n",
    "opt.n_blocks_global = 9  # Number of residual blocks in the generator.\n",
    "opt.n_blocks_local = 3  # Number of residual blocks in the generator.\n",
    "opt.n_local_enhancers = 1  # Number of local enhancers in the generator.\n",
    "\n",
    "# Define the parameters for the Discriminators.\n",
    "opt.num_D = 3  # Number of discriminators.\n",
    "opt.n_layers_D = 3  # Number of layers in the discriminator.\n",
    "opt.ndf = 32  # Number of filters in the discriminator.\n",
    "\n",
    "# Define general training parameters.\n",
    "opt.gpu_ids = [0] # GPU ids to use.\n",
    "opt.norm = \"instance\"  # Normalization layer in the generator.\n",
    "opt.use_dropout = \"\"  # Use dropout in the generator (fixed at 0.2).\n",
    "opt.batchSize = 8  # Batch size.\n",
    "\n",
    "# Create a visualizer to perform image processing and visualization\n",
    "visualizer = Visualizer(opt)\n",
    "\n",
    "# Here will first start training a model from scrach however we can continue to train from a previously trained model by setting the following parameters.\n",
    "opt.continue_train = False\n",
    "if opt.continue_train:\n",
    "    iter_path = os.path.join(opt.checkpoints_dir, opt.name, \"iter.txt\")\n",
    "    try:\n",
    "        start_epoch, epoch_iter = np.loadtxt(iter_path, delimiter=\",\", dtype=int)\n",
    "    except:\n",
    "        start_epoch, epoch_iter = 1, 0\n",
    "    print(\"Resuming from epoch %d at iteration %d\" % (start_epoch, epoch_iter))\n",
    "else:\n",
    "    start_epoch, epoch_iter = 1, 0\n",
    "    \n",
    "print('------------ Options -------------')\n",
    "for k, v in sorted(vars(opt).items()):\n",
    "    print('%s: %s' % (str(k), str(v)))\n",
    "print('-------------- End ----------------')\n",
    "\n",
    "# Set the number of epoch to be 1 for demonstration purposes\n",
    "opt.n_epochs = 2\n",
    "# Initialize the model\n",
    "phase2nuclei_model = create_model(opt)\n",
    "# Define Optimizers for G and D\n",
    "optimizer_G, optimizer_D = (\n",
    "    phase2nuclei_model.module.optimizer_G,\n",
    "    phase2nuclei_model.module.optimizer_D,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "beb25970",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [],
   "source": [
    "train_model(\n",
    "    opt,\n",
    "    phase2nuclei_model,\n",
    "    visualizer,\n",
    "    dataset_train,\n",
    "    dataset_val,\n",
    "    optimizer_G,\n",
    "    optimizer_D,\n",
    "    start_epoch,\n",
    "    epoch_iter,\n",
    "    writer,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "83f2338a",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 0
   },
   "source": [
    "<div class=\"alert alert-info\">\n",
    "\n",
    "## A heads up of what to expect from the training...\n",
    "<br><br>\n",
    "**- Visualise results**: We can observe how the performance improves over time using the images tab and the sliding window.\n",
    "<br><br>\n",
    "**- Discriminator Predicted Probabilities**: We plot the discriminator's predicted probabilities that the phase with fluorescence is phase and fluorescence and that the phase with virtual stain is phase with virtual stain. It is typically trained until the discriminator can no longer classify whether or not the generated images are real or fake better than a random guess (p(0.5)). We plot this for both the training and validation datasets.\n",
    "<br><br>\n",
    "**- Adversarial Loss**: We can formulate the adversarial loss as a Least Squared Error Loss in which for real data the discriminator should output a value close to 1 and for fake data a value close to 0. The generator's goal is to make the discriminator output a value as close to 1 for fake data. We plot the least squared error loss.\n",
    "<br><br>\n",
    "**- Feature Matching Loss**: Both networks are also trained using the generator feature matching loss which encourages the generator to produce images that contain similar statistics to the real images at each scale. We also plot the feature matching L1 loss for the training and validation sets together to observe the performance and how the model is fitting the data.\n",
    "<br><br>\n",
    "This implementation allows for the turning on/off of the least-square loss term by setting the opt.no_lsgan flag to the model options. As well as the turning off of the feature matching loss term by setting the opt.no_ganFeat_loss flag to the model options. Something you might want to explore in the next section!<br><br>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e94c94db",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 0
   },
   "source": [
    "<div class=\"alert alert-success\">\n",
    "    \n",
    "## Checkpoint 1\n",
    "\n",
    "Congratulations! You should now have a better understanding of how to train a Pix2PixHD GAN model for translating from phase to nuclei. You should also have a good understanding of the different loss components of Pix2PixHD GAN and how they are used to train the model.\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d13cbb1a",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 0,
    "title": "<a ></a>"
   },
   "source": [
    "# Part 2: Load & Assess trained Pix2PixGAN using tensorboard, discuss performance of the model.\n",
    "--------------------------------------------------\n",
    "Learning goals:\n",
    "- Load a pre-trained Pix2PixHD GAN model for either phase to nuclei.\n",
    "- Discuss the loss components of Pix2PixHD GAN and how they are used to train the model.\n",
    "- Evaluate the fit of the model on the train and validation datasets.\n",
    "\n",
    "In this part, we will evaluate the performance of the pre-trained model. We will begin by looking qualitatively at the model predictions, then dive into the different loss plots. We will discuss the implications of different hyper-parameter combinations for the performance of the model.\n",
    "\n",
    "If you are having issues loading the tensorboard session click \"Launch TensorBoard session\". You should then be able to add the log_dir path below and a tensorboard session shouls then load."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99c8ab80",
   "metadata": {},
   "outputs": [],
   "source": [
    "log_dir = f\"{top_dir}/model_tensorboard/{opt.name}/\"\n",
    "%reload_ext tensorboard\n",
    "%tensorboard --logdir $log_dir"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a021ecc",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "## Training Results\n",
    "Please note down your thoughts about the following questions...\n",
    "\n",
    "**- What do you notice about the virtual staining predictions? How do they appear compared to the regression-based approach? Can you spot any hallucinations?**<br><br> \n",
    "**- What do you notice about the probabilities of the discriminators? How do the values compare during training compared to validation?**<br><br>\n",
    "**- What do you notice about the feature matching L1 loss?**<br><br>\n",
    "**- What do you notice about the least-square loss?**<br><br>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa14c1ce",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "<div class=\"alert alert-success\">\n",
    "    \n",
    "## Checkpoint 2\n",
    "Congratulations! You should now have a better understanding the different loss components of Pix2PixHD GAN and how they are used to train the model. You should also have a good understanding of the fit of the model during training on the training and validation datasets.\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db973adc",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "# Part 3: Evaluate performance of the virtual staining on unseen data.\n",
    "--------------------------------------------------\n",
    "## Evaluate the performance of the model.\n",
    "We now look at the same metrics of performance of the previous model. We typically evaluate the model performance on a held out test data. \n",
    "\n",
    "Steps:\n",
    "- Define our model parameters for the pre-trained model (these are the same parameters as shown in earlier cells but copied here for clarity).\n",
    "- Load the test data.\n",
    "\n",
    "We will first load the test data using the same format as the training and validation data. We will then use the model to sample a virtual nuclei staining soltuion from the phase image. We will then evaluate the performance of the model using the following metrics:\n",
    "\n",
    "Pixel-level metrics:\n",
    "- [Peak-Signal-to-Noise-Ratio (PSNR)](https://en.wikipedia.org/wiki/Peak_signal-to-noise_ratio).\n",
    "- [Structural Similarity Index Measure (SSIM)](https://en.wikipedia.org/wiki/Structural_similarity).\n",
    "- [Pearson Correlation Coefficient (PCC)](https://en.wikipedia.org/wiki/Pearson_correlation_coefficient).\n",
    "\n",
    "Instance-level metrics via [Cellpose masks](https://cellpose.org/):\n",
    "- [Accuracy](https://en.wikipedia.org/wiki/Accuracy_and_precision#In_binary_classification)\n",
    "- [Jaccard Index](https://en.wikipedia.org/wiki/Jaccard_index)\n",
    "- [Dice Score](https://en.wikipedia.org/wiki/S%C3%B8rensen%E2%80%93Dice_coefficient)\n",
    "- [Mean Average Precision](https://en.wikipedia.org/wiki/Evaluation_measures_(information_retrieval)#Mean_average_precision)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7106aaad",
   "metadata": {},
   "outputs": [],
   "source": [
    "opt = TestOptions().parse(save=False)\n",
    "\n",
    "# Define the parameters for the dataset.\n",
    "opt.dataroot = output_image_folder\n",
    "opt.data_type = 16  # Data type of the images.\n",
    "opt.loadSize = 512  # Size of the loaded phase image.\n",
    "opt.input_nc = 1  # Number of input channels.\n",
    "opt.output_nc = 1  # Number of output channels.\n",
    "opt.target = \"nuclei\"  # \"nuclei\" or \"cyto\" depending on your choice of target for virtual stain.\n",
    "opt.resize_or_crop = \"none\"  # Scaling and cropping of images at load time [resize_and_crop|crop|scale_width|scale_width_and_crop|none].\n",
    "opt.batchSize = 1 # Batch size for training\n",
    "\n",
    "# Define the model parameters for the pre-trained model.\n",
    "\n",
    "# Define the parameters for the Generator.\n",
    "opt.ngf = 64  # Number of filters in the generator.\n",
    "opt.n_downsample_global = 4  # Number of downsampling layers in the generator.\n",
    "opt.n_blocks_global = 9  # Number of residual blocks in the generator.\n",
    "opt.n_blocks_local = 3  # Number of residual blocks in the generator.\n",
    "opt.n_local_enhancers = 1  # Number of local enhancers in the generator.\n",
    "\n",
    "# Define the parameters for the Discriminators.\n",
    "opt.num_D = 3  # Number of discriminators.\n",
    "opt.n_layers_D = 3  # Number of layers in the discriminator.\n",
    "opt.ndf = 32  # Number of filters in the discriminator.\n",
    "\n",
    "# Define general training parameters.\n",
    "opt.gpu_ids= [0]  # GPU ids to use.\n",
    "opt.norm = \"instance\"  # Normalization layer in the generator.\n",
    "opt.use_dropout = \"\"  # Use dropout in the generator (fixed at 0.2).\n",
    "opt.batchSize = 8  # Batch size.\n",
    "\n",
    "# Define loss functions.\n",
    "opt.no_vgg_loss = \"\"  # Turn off VGG loss\n",
    "opt.no_ganFeat_loss = \"\"  # Turn off feature matching loss\n",
    "opt.no_lsgan = \"\"  # Turn off least square loss\n",
    "\n",
    "# Additional Inference parameters\n",
    "opt.name = f\"dlmbl_vsnuclei\"\n",
    "opt.how_many = 112  # Number of images to generate.\n",
    "opt.checkpoints_dir = f\"{top_dir}/model_weights/\"  # Path to the model checkpoints.\n",
    "opt.results_dir = f\"{top_dir}/GAN_code/GANs_MI2I/pre_trained/{opt.name}/inference_results/\"  # Path to store the results.\n",
    "opt.which_epoch = \"latest\"  # or specify the epoch number \"40\"\n",
    "opt.phase = \"test\"\n",
    "\n",
    "opt.nThreads = 1  # test code only supports nThreads = 1\n",
    "opt.batchSize = 1  # test code only supports batchSize = 1\n",
    "opt.serial_batches = True  # no shuffle\n",
    "opt.no_flip = True  # no flip\n",
    "Path(opt.results_dir).mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "# Load the test data.\n",
    "test_data_loader = CreateDataLoader(opt)\n",
    "test_dataset = test_data_loader.load_data()\n",
    "visualizer = Visualizer(opt)\n",
    "print(f\"Total Test Images = {len(test_data_loader)}\")\n",
    "# Load pre-trained model\n",
    "model = create_model(opt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b29abc8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate & save predictions in the results directory.\n",
    "inference_model(test_dataset, opt, model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54514ba9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Gather results for evaluation\n",
    "virtual_stain_paths = sorted([i for i in Path(opt.results_dir).glob(\"**/*.tiff\")])\n",
    "target_stain_paths = sorted([i for i in Path(f\"{output_image_folder}/{translation_task}/test/\").glob(\"**/*.tiff\")])\n",
    "phase_paths = sorted([i for i in Path(f\"{output_image_folder}/input/test/\").glob(\"**/*.tiff\")])\n",
    "assert (len(virtual_stain_paths) == len(target_stain_paths) == len(phase_paths)\n",
    "), f\"Number of images do not match. {len(virtual_stain_paths)},{len(target_stain_paths)} {len(phase_paths)} \"\n",
    "\n",
    "# Create arrays to store the images.\n",
    "virtual_stains = np.zeros((len(virtual_stain_paths), 512, 512))\n",
    "target_stains = virtual_stains.copy()\n",
    "phase_images = virtual_stains.copy()\n",
    "# Load the images and store them in the arrays.\n",
    "for index, (v_path, t_path, p_path) in tqdm(\n",
    "    enumerate(zip(virtual_stain_paths, target_stain_paths, phase_paths))\n",
    "):\n",
    "    virtual_stain = imread(v_path)\n",
    "    phase_image = imread(p_path)\n",
    "    target_stain = imread(t_path)\n",
    "    # Append the images to the arrays.\n",
    "    phase_images[index] = phase_image\n",
    "    target_stains[index] = target_stain\n",
    "    virtual_stains[index] = virtual_stain"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0dbb8ea0",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 0,
    "tags": []
   },
   "source": [
    "<div class=\"alert alert-info\">\n",
    "\n",
    "### Task 3.1 Visualise the results of the model on the test set.\n",
    "\n",
    "Create a matplotlib plot that visalises random samples of the phase images, target stains, and virtual stains.\n",
    "If you can incorporate the crop function below to zoom in on the images that would be great!\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb768fe1",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "# Define a function to crop the images so we can zoom in.\n",
    "def crop(img, crop_size, loc='center'):\n",
    "    \"\"\"\n",
    "    Crop the input image.\n",
    "\n",
    "    Parameters:\n",
    "    img (ndarray): The image to be cropped.\n",
    "    crop_size (int): The size of the crop.\n",
    "    loc (str): The type of crop to perform. Can be 'center' or 'random'.\n",
    "\n",
    "    Returns:\n",
    "    ndarray: The cropped image array.\n",
    "    \"\"\"\n",
    "    # Dimension of input array\n",
    "    width, height = img.shape\n",
    "\n",
    "    # Generate random coordinates for the crop\n",
    "    max_y = height - crop_size\n",
    "    max_x = max_y\n",
    "\n",
    "    if loc == 'random':\n",
    "        start_y = np.random.randint(0, max_y + 1)\n",
    "        start_x = np.random.randint(0, max_x + 1)\n",
    "        end_y = start_y + crop_size\n",
    "        end_x = start_x + crop_size\n",
    "    elif loc == 'center':\n",
    "        start_x = (width - crop_size) // 2\n",
    "        start_y = (height - crop_size) // 2\n",
    "        end_y = height - start_y\n",
    "        end_x = width - start_x\n",
    "    else:\n",
    "        raise ValueError(f'Unknown crop type {loc}')\n",
    "\n",
    "    # Crop array using slicing\n",
    "    crop_array = img[start_x:end_x, start_y:end_y]\n",
    "    return crop_array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "133ab59e",
   "metadata": {
    "tags": [
     "task"
    ]
   },
   "outputs": [],
   "source": [
    "##########################\n",
    "######## TODO ########\n",
    "##########################\n",
    "\n",
    "def visualise_results():\n",
    "    # Your code here\n",
    "    pass"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "91ea20ac",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 0,
    "tags": []
   },
   "source": [
    "<div class=\"alert alert-info\">\n",
    "\n",
    "### Task 3.2 Compute pixel-level metrics\n",
    "\n",
    "Compute the pixel-level metrics for the virtual stains and target stains.\n",
    "\n",
    "The following code will compute the following:\n",
    "- the pixel-based metrics  (Pearson correlation, SSIM, PSNR) for the virtual stains and target stains.\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "baf66869",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [],
   "source": [
    "\n",
    "# Define the function to perform  minmax normalization which is required for the pixel-level metrics.\n",
    "def min_max_scale(input):\n",
    "    return (input - np.min(input)) / (np.max(input) - np.min(input))\n",
    "\n",
    "# Create a dataframe to store the pixel-level metrics.\n",
    "test_pixel_metrics = pd.DataFrame(\n",
    "    columns=[\"model\", \"fov\",\"pearson_nuc\", \"ssim_nuc\", \"psnr_nuc\"]\n",
    ")\n",
    "\n",
    "# Compute the pixel-level metrics.\n",
    "for i, (target_stain, predicted_stain) in tqdm(enumerate(zip(target_stains, virtual_stains))):\n",
    "    fov = str(virtual_stain_paths[i]).split(\"/\")[-1].split(\".\")[0]\n",
    "    minmax_norm_target = min_max_scale(target_stain)\n",
    "    minmax_norm_predicted = min_max_scale(predicted_stain)\n",
    "    \n",
    "    # Compute SSIM\n",
    "    ssim_nuc = metrics.structural_similarity(\n",
    "        minmax_norm_target, minmax_norm_predicted, data_range=1\n",
    "    )\n",
    "    # Compute Pearson correlation\n",
    "    pearson_nuc = np.corrcoef(\n",
    "        minmax_norm_target.flatten(), minmax_norm_predicted.flatten()\n",
    "        )[0, 1]\n",
    "    # Compute PSNR\n",
    "    psnr_nuc = metrics.peak_signal_noise_ratio(\n",
    "        minmax_norm_target, minmax_norm_predicted, data_range=1\n",
    "    )\n",
    "    \n",
    "    test_pixel_metrics.loc[len(test_pixel_metrics)] = {\n",
    "            \"model\": \"pix2pixHD\",\n",
    "            \"fov\":fov,\n",
    "            \"pearson_nuc\": pearson_nuc,\n",
    "            \"ssim_nuc\": ssim_nuc,\n",
    "            \"psnr_nuc\": psnr_nuc, \n",
    "        }\n",
    "    \n",
    "test_pixel_metrics.boxplot(\n",
    "    column=[\"pearson_nuc\", \"ssim_nuc\"],\n",
    "    rot=30,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "721e9b16",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [],
   "source": [
    "test_pixel_metrics.boxplot(\n",
    "    column=[\"psnr_nuc\"],\n",
    "    rot=30,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "abb07bea",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [],
   "source": [
    "test_pixel_metrics.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46e19877",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 0
   },
   "source": [
    "## Inference Pixel-level Results\n",
    "Please note down your thoughts about the following questions...\n",
    "\n",
    "- What do these metrics tells us about the performance of the model?\n",
    "- How do the pixel-level metrics compare to the regression-based approach?\n",
    "- Could these metrics be skewed by the presence of hallucinations or background pilxels in the virtual stains?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b6e7b6b",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 0
   },
   "source": [
    "<div class=\"alert alert-info\">\n",
    "\n",
    "### Task 3.3 Compute instance-level metrics\n",
    "\n",
    "- Compute the instance-level metrics for the virtual stains and target stains.\n",
    "- Instance metrics include the accuracy (average correct predictions with 0.5 threshold), jaccard index (intersection over union (IoU)) dice score (2x intersection over union), mean average precision, mean average precision at 50% IoU, mean average precision at 75% IoU, and mean average recall at 100% IoU.\n",
    "\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d720c0f1",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [],
   "source": [
    "\n",
    "# Use the same function as previous part to extract the nuclei masks from pre-trained cellpose model.\n",
    "def cellpose_segmentation(prediction:ArrayLike,target:ArrayLike)->Tuple[torch.ShortTensor]:\n",
    "    # NOTE these are hardcoded for this notebook and A549 dataset\n",
    "    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "    cp_nuc_kwargs = {\n",
    "        \"diameter\": 65,\n",
    "        \"cellprob_threshold\": 0.0, \n",
    "    }\n",
    "    cellpose_model = models.CellposeModel(\n",
    "            gpu=True, model_type='nuclei', device=torch.device(device)\n",
    "    )\n",
    "    pred_label, _, _ = cellpose_model.eval(prediction, **cp_nuc_kwargs)\n",
    "    target_label, _, _ = cellpose_model.eval(target, **cp_nuc_kwargs)\n",
    "\n",
    "    pred_label = pred_label.astype(np.int32)\n",
    "    target_label = target_label.astype(np.int32)\n",
    "    pred_label = torch.ShortTensor(pred_label)\n",
    "    target_label = torch.ShortTensor(target_label)\n",
    "\n",
    "    return (pred_label,target_label)\n",
    "\n",
    "# Define dataframe to store the segmentation metrics.\n",
    "test_segmentation_metrics= pd.DataFrame(\n",
    "    columns=[\"model\", \"fov\",\"masks_per_fov\",\"accuracy\",\"dice\",\"jaccard\",\"mAP\",\"mAP_50\",\"mAP_75\",\"mAR_100\"]\n",
    ")\n",
    "# Define tuple to store the segmentation results. Each value in the tuple is a dictionary containing the model name, fov, predicted label, predicted stain, target label, and target stain.\n",
    "segmentation_results = ()\n",
    "\n",
    "for i, (target_stain, predicted_stain) in tqdm(enumerate(zip(target_stains, virtual_stains))):\n",
    "    fov = str(virtual_stain_paths)[i].spilt(\"/\")[-1].split(\".\")[0]\n",
    "    minmax_norm_target = min_max_scale(target_stain)\n",
    "    minmax_norm_predicted = min_max_scale(predicted_stain)\n",
    "    # Compute the segmentation masks.\n",
    "    pred_label, target_label = cellpose_segmentation(minmax_norm_predicted, minmax_norm_target)\n",
    "     # Binary labels\n",
    "    pred_label_binary = pred_label > 0\n",
    "    target_label_binary = target_label > 0\n",
    "\n",
    "    # Use Coco metrics to get mean average precision\n",
    "    coco_metrics = mean_average_precision(pred_label, target_label)\n",
    "    # Find unique number of labels\n",
    "    num_masks_fov = len(np.unique(pred_label))\n",
    "    # Find unique number of labels\n",
    "    num_masks_fov = len(np.unique(pred_label))\n",
    "    # Compute the segmentation metrics.\n",
    "    test_segmentation_metrics.loc[len(test_segmentation_metrics)] = {\n",
    "        \"model\": \"pix2pixHD\",\n",
    "        \"fov\":fov,\n",
    "        \"masks_per_fov\": num_masks_fov,\n",
    "        \"accuracy\": accuracy(pred_label_binary, target_label_binary, task=\"binary\").item(),\n",
    "        \"dice\":  dice(pred_label_binary, target_label_binary).item(),\n",
    "        \"jaccard\": jaccard_index(pred_label_binary, target_label_binary, task=\"binary\").item(),\n",
    "        \"mAP\":coco_metrics[\"map\"].item(),\n",
    "        \"mAP_50\":coco_metrics[\"map_50\"].item(),\n",
    "        \"mAP_75\":coco_metrics[\"map_75\"].item(),\n",
    "        \"mAR_100\":coco_metrics[\"mar_100\"].item()\n",
    "        }\n",
    "    # Store the segmentation results.\n",
    "    segmentation_result = {\n",
    "        \"model\": \"pix2pixHD\",\n",
    "        \"fov\":fov,\n",
    "        \"phase_image\": phase_images[i],\n",
    "        \"pred_label\": pred_label,\n",
    "        \"pred_stain\": predicted_stain,\n",
    "        \"target_label\": target_label,\n",
    "        \"target_stain\": target_stain,\n",
    "    }\n",
    "    segmentation_results += (segmentation_result,)\n",
    "\n",
    "test_segmentation_metrics.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77730fbb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define function to visualize the segmentation results.\n",
    "def visualise_results_and_masks(segmentation_results: Tuple[dict], segmentation_metrics: pd.DataFrame, rows: int = 5, crop_size: int = None, crop_type: str = 'center'):\n",
    "\n",
    "    # Sample a subset of the segmentation results.\n",
    "    sample_indices = np.random.choice(len(phase_images),rows)\n",
    "    print(sample_indices)\n",
    "    segmentation_metrics = segmentation_metrics.iloc[sample_indices,:]\n",
    "    segmentation_results = [segmentation_results[i] for i in sample_indices]\n",
    "    # Define the figure and axes.\n",
    "    fig, axes = plt.subplots(rows, 5, figsize=(rows*3, 15))\n",
    "\n",
    "    # Visualize the segmentation results.\n",
    "    for i in range(len((segmentation_results))):\n",
    "        segmentation_metric = segmentation_metrics.iloc[i]\n",
    "        result = segmentation_results[i]\n",
    "        phase_image = result[\"phase_image\"]\n",
    "        target_stain = result[\"target_stain\"]\n",
    "        target_label = result[\"target_label\"]\n",
    "        pred_stain = result[\"pred_stain\"]\n",
    "        pred_label = result[\"pred_label\"]\n",
    "        # Crop the images if required. Zoom into instances\n",
    "        if crop_size is not None:\n",
    "            phase_image = crop(phase_image, crop_size, crop_type)\n",
    "            target_stain = crop(target_stain, crop_size, crop_type)\n",
    "            target_label = crop(target_label, crop_size, crop_type)\n",
    "            pred_stain = crop(pred_stain, crop_size, crop_type)\n",
    "            pred_label = crop(pred_label, crop_size, crop_type)\n",
    "        \n",
    "        axes[i, 0].imshow(phase_image, cmap=\"gray\")\n",
    "        axes[i, 0].set_title(\"Phase\")\n",
    "        axes[i, 1].imshow(\n",
    "            target_stain,\n",
    "            cmap=\"gray\",\n",
    "            vmin=np.percentile(target_stain, 1),\n",
    "            vmax=np.percentile(target_stain, 99),\n",
    "        )\n",
    "        axes[i, 1].set_title(\"Target Fluorescence\")\n",
    "        axes[i, 2].imshow(pred_stain, cmap=\"gray\")\n",
    "        axes[i, 2].set_title(\"Virtual Stain\")\n",
    "        axes[i, 3].imshow(target_label, cmap=\"inferno\")\n",
    "        axes[i, 3].set_title(\"Target Fluorescence Mask\")\n",
    "        axes[i, 4].imshow(pred_label, cmap=\"inferno\")\n",
    "        # Add Metric values to the title\n",
    "        axes[i, 4].set_title(f\"Virtual Stain Mask\\nAcc:{segmentation_metric['accuracy']:.2f} Dice:{segmentation_metric['dice']:.2f}\\nJaccard:{segmentation_metric['jaccard']:.2f} MAP:{segmentation_metric['mAP']:.2f}\")\n",
    "    # Turn off the axes.\n",
    "    for ax in axes.flatten():\n",
    "        ax.axis(\"off\")\n",
    "\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "    \n",
    "visualise_results_and_masks(segmentation_results,test_segmentation_metrics, crop_size=256, crop_type='center')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "692679b6",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "685d7e60",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 0
   },
   "source": [
    "## Inference Instance-level Results\n",
    "Please note down your thoughts about the following questions...\n",
    "\n",
    "- What do these metrics tells us about the performance of the model?\n",
    "- How does the performance compare to when looking at pixel-level metrics?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b75f675d",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "<div class=\"alert alert-success\">\n",
    "    \n",
    "## Checkpoint 3\n",
    "\n",
    "Congratulations! You have generated predictions from a pre-trained model and evaluated the performance of the model on unseen data. You have computed pixel-level metrics and instance-level metrics to evaluate the performance of the model. You may have also began training your own Pix2PixHD GAN models with alternative hyperparameters.\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c1c584e",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 0
   },
   "source": [
    "# Part 4. Visualise Regression vs Generative Modelling Approaches\n",
    "--------------------------------------------------"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae61bb14",
   "metadata": {
    "tags": [
     "task"
    ]
   },
   "outputs": [],
   "source": [
    "# Load Viscy Virtual Stains\n",
    "viscy_results_path = \"/ADD/PATH/TO/RESULTS/HERE\"\n",
    "viscy_stain_paths = sorted([i for i in Path(viscy_results_path).glob(\"**/*.tiff\")])\n",
    "assert len(viscy_stain_paths) == len(virtual_stain_paths), \"Number of images do not match.\"\n",
    "visy_stains = np.zeros((len(viscy_stain_paths), 512, 512))\n",
    "for index, v_path in enumerate(viscy_stain_paths):\n",
    "    viscy_stain = imread(v_path)\n",
    "    visy_stains[index] = viscy_stain\n",
    "\n",
    "##########################\n",
    "######## TODO ########\n",
    "##########################\n",
    "\n",
    "\n",
    "def visualise_both_methods():\n",
    "    # Your code here\n",
    "    pass"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33ee4c8b",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 0
   },
   "source": [
    "<div class=\"alert alert-success\">\n",
    "    \n",
    "## Checkpoint 4\n",
    "\n",
    "Congratulations! You should now have a better understanding of the difference in performance for image translation when approaching the problem using a regression vs. generative modelling approaches!\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a8b4cb73",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 0
   },
   "source": [
    "# Part 5: BONUS: Sample different virtual staining solutions from the GAN using MC-Dropout and explore the uncertainty in the virtual stain predictions.\n",
    "--------------------------------------------------\n",
    "Steps:\n",
    "- Load the pre-trained model.\n",
    "- Generate multiple predictions for the same input image.\n",
    "- Compute the pixel-wise variance across the predictions.\n",
    "- Visualise the pixel-wise variance to explore the uncertainty in the virtual stain predictions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db38f032",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use the same model and dataloaders as before.\n",
    "# Load the test data.\n",
    "test_data_loader = CreateDataLoader(opt)\n",
    "test_dataset = test_data_loader.load_data()\n",
    "visualizer = Visualizer(opt)\n",
    "\n",
    "# Load pre-trained model\n",
    "opt.variational_inf_runs = 100 # Number of samples per phase input\n",
    "opt.variation_inf_path = f\"./GAN_code/GANs_MI2I/pre_trained/{opt.name}/samples/\"  # Path to store the samples.\n",
    "opt.results_dir = f\"{top_dir}/GAN_code/GANs_MI2I/pre_trained/dlmbl_vsnuclei/sampling_results\"\n",
    "opt.dropout_variation_inf = True  # Use dropout during inference.\n",
    "model = create_model(opt)\n",
    "# Generate & save predictions in the variation_inf_path directory.\n",
    "sampling(test_dataset, opt, model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a94162a7",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "# Visualise Samples                                      \n",
    "samples = sorted([i for i in Path(f\"{opt.results_dir}\").glob(\"**/*.tif*\")])\n",
    "assert len(samples) == 5\n",
    "# Create arrays to store the images.\n",
    "sample_images = np.zeros((5,100, 512, 512)) # (samples, images, height, width)\n",
    "# Load the images and store them in the arrays.\n",
    "for index, sample_path in tqdm(enumerate(samples)):\n",
    "    sample_image = imread(sample_path)\n",
    "    # Append the images to the arrays.\n",
    "    sample_images[index] = sample_image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "efa01bad",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a matplotlib plot with animation through images.\n",
    "def animate_images(images):\n",
    "    # Expecting images to have shape (frames, height, width)\n",
    "    fig, ax = plt.subplots()\n",
    "    ax.axis('off')\n",
    "    \n",
    "    # Make sure images are in (frames, height, width) order\n",
    "    images = images.transpose(0, 2, 1) if images.shape[1] == images.shape[2] else images\n",
    "    \n",
    "    imgs = []\n",
    "    for i in range(min(100, len(images))):  # Ensure you don't exceed the number of frames\n",
    "        im = ax.imshow(images[i], animated=True)\n",
    "        imgs.append([im])\n",
    "\n",
    "    ani = animation.ArtistAnimation(fig, imgs, interval=100, blit=False, repeat_delay=1000)\n",
    "    \n",
    "    # Display the animation\n",
    "    # plt.close(fig)\n",
    "    display(HTML(ani.to_jshtml()))\n",
    "\n",
    "# Example call with sample_images[0]\n",
    "animate_images(sample_images[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "670b504f",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "<div class=\"alert alert-success\">\n",
    "    \n",
    "## Checkpoint 5\n",
    "\n",
    "Congratulations! This is the end of the conditional generative modelling approach to image translation notebook. You have trained and examined the loss components of Pix2PixHD GAN. You have compared the results of a regression-based approach vs. generative modelling approach and explored the variability in virtual staining solutions. I hope you have enjoyed learning experience!\n",
    "</div>"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "cell_metadata_filter": "all",
   "main_language": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}