solution.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "09d53636",
   "metadata": {
    "lines_to_next_cell": 0,
    "tags": []
   },
   "source": [
    "# Image translation (Virtual Staining)\n",
    "\n",
    "Written by Eduardo Hirata-Miyasaki, Ziwen Liu, and Shalin Mehta, CZ Biohub San Francisco\n",
    "\n",
    "## Overview\n",
    "\n",
    "In this exercise, we will _virtually stain_ the nuclei and plasma membrane from the quantitative phase image (QPI), i.e., translate QPI images into fluoresence images of nuclei and plasma membranes.\n",
    "QPI encodes multiple cellular structures and virtual staining decomposes these structures. After the model is trained, one only needs to acquire label-free QPI data.\n",
    "This strategy solves the problem as \"multi-spectral imaging\", but is more compatible with live cell imaging and high-throughput screening.\n",
    "Virtual staining is often a step towards multiple downstream analyses: segmentation, tracking, and cell state phenotyping.\n",
    "\n",
    "In this exercise, you will:\n",
    "- Train a model to predict the fluorescence images of nuclei and plasma membranes from QPI images\n",
    "- Make it robust to variations in imaging conditions using data augmentions\n",
    "- Segment the cells\n",
    "- Use regression and segmentation metrics to evalute the models\n",
    "- Visualize the image transform learned by the model\n",
    "- Understand the failure modes of the trained model\n",
    "\n",
    "[![HEK293T](https://raw.githubusercontent.com/mehta-lab/VisCy/main/docs/figures/svideo_1.png)](https://github.com/mehta-lab/VisCy/assets/67518483/d53a81eb-eb37-44f3-b522-8bd7bddc7755)\n",
    "(Click on image to play video)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c351b9ee",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Goals\n",
    "\n",
    "#### Part 1: Train a virtual staining model\n",
    "\n",
    "  - Explore OME-Zarr using [iohub](https://czbiohub-sf.github.io/iohub/main/index.html)\n",
    "  and the high-content-screen (HCS) format.\n",
    "  - Use our `viscy.data.HCSDataloader()` dataloader and explore the  3 channel (phase, fluoresecence nuclei and cell membrane) \n",
    "  A549 cell dataset. \n",
    "  - Implement data augmentations [MONAI](https://monai.io/) to train a robust model to imaging parameters and conditions. \n",
    "  - Use tensorboard to log the augmentations, training and validation losses and batches\n",
    "  - Start the training of the UNeXt2 model to predict nuclei and membrane from phase images.\n",
    "\n",
    "#### Part 2:Evaluate the model to translate phase into fluorescence.\n",
    "  - Compare the performance of your trained model with the _VSCyto2D_ pre-trained model.\n",
    "  - Evaluate the model using pixel-level and instance-level metrics.\n",
    "\n",
    "#### Part 3: Visualize the image transforms learned by the model and explore the model's regime of validity\n",
    "  - Visualize the first 3 principal componets mapped to a color space in each encoder and decoder block.\n",
    "  - Explore the model's regime of validity by applying blurring and scaling transforms to the input phase image.\n",
    "\n",
    "#### For more information:\n",
    "Checkout [VisCy](https://github.com/mehta-lab/VisCy),\n",
    "our deep learning pipeline for training and deploying computer vision models\n",
    "for image-based phenotyping including the robust virtual staining of landmark organelles.\n",
    "\n",
    "VisCy exploits recent advances in data and metadata formats\n",
    "([OME-zarr](https://www.nature.com/articles/s41592-021-01326-w)) and DL frameworks,\n",
    "[PyTorch Lightning](https://lightning.ai/) and [MONAI](https://monai.io/).\n",
    "\n",
    "### References\n",
    "- [Liu, Z. and Hirata-Miyasaki, E. et al. (2024) Robust Virtual Staining of Cellular Landmarks](https://www.biorxiv.org/content/10.1101/2024.05.31.596901v2.full.pdf)\n",
    "- [Guo et al. (2020) Revealing architectural order with quantitative label-free imaging and deep learning. eLife](https://elifesciences.org/articles/55502)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "deb2f467",
   "metadata": {
    "tags": []
   },
   "source": [
    "<div class=\"alert alert-success\">\n",
    "The exercise is organized in 3 parts:\n",
    "\n",
    "<ul>\n",
    "<li><b>Part 1</b> - Train a virtual staining model using iohub (I/O library), VisCy dataloaders, and tensorboard</li>\n",
    "<li><b>Part 2</b> - Evaluate the model to translate phase into fluorescence.</li>\n",
    "<li><b>Part 3</b> - Visualize the image transforms learned by the model and explore the model's regime of validity.</li>\n",
    "</ul>\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a7ea8019",
   "metadata": {
    "lines_to_next_cell": 0,
    "tags": []
   },
   "source": [
    "<div class=\"alert alert-danger\">\n",
    "Set your python kernel to <span style=\"color:black;\">06_image_translation</span>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f36840b9",
   "metadata": {},
   "source": [
    "# Part 1: Log training data to tensorboard, start training a model.\n",
    "---------\n",
    "Learning goals:\n",
    "\n",
    "- Load the OME-zarr dataset and examine the channels (A549).\n",
    "- Configure and understand the data loader.\n",
    "- Log some patches to tensorboard.\n",
    "- Initialize a 2D UNeXt2 model for virtual staining of nuclei and membrane from phase.\n",
    "- Start training the model to predict nuclei and membrane from phase."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f01b54ab",
   "metadata": {
    "title": "Imports"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "from glob import glob\n",
    "from pathlib import Path\n",
    "from typing import Tuple\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import torch\n",
    "import torchview\n",
    "import torchvision\n",
    "from cellpose import models\n",
    "from iohub import open_ome_zarr\n",
    "from iohub.reader import print_info\n",
    "from lightning.pytorch import seed_everything\n",
    "from lightning.pytorch.loggers import TensorBoardLogger\n",
    "from natsort import natsorted\n",
    "from numpy.typing import ArrayLike\n",
    "from skimage import metrics  # for metrics.\n",
    "# pytorch lightning wrapper for Tensorboard.\n",
    "from skimage.color import label2rgb\n",
    "from torch.utils.tensorboard import SummaryWriter  # for logging to tensorboard\n",
    "from torchmetrics.functional import accuracy, dice, jaccard_index\n",
    "from tqdm import tqdm\n",
    "# HCSDataModule makes it easy to load data during training.\n",
    "from viscy.data.hcs import HCSDataModule\n",
    "from viscy.evaluation.evaluation_metrics import mean_average_precision\n",
    "# Trainer class and UNet.\n",
    "from viscy.light.engine import MixedLoss, VSUNet\n",
    "from viscy.light.trainer import VSTrainer\n",
    "# training augmentations\n",
    "from viscy.transforms import (NormalizeSampled, RandAdjustContrastd,\n",
    "                              RandAffined, RandGaussianNoised,\n",
    "                              RandGaussianSmoothd, RandScaleIntensityd,\n",
    "                              RandWeightedCropd)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25c9e554",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "# seed random number generators for reproducibility.\n",
    "seed_everything(42, workers=True)\n",
    "\n",
    "# Paths to data and log directory\n",
    "top_dir = Path(\n",
    "    \"/mnt/efs/dlmbl/share/\"\n",
    ")  # If this fails, make sure this to point to your data directory in the shared mounting point inside /dlmbl/data\n",
    "\n",
    "# Path to the training data\n",
    "data_path = (\n",
    "    top_dir / \"06_image_translation/training/a549_hoechst_cellmask_train_val.zarr\"\n",
    ")\n",
    "\n",
    "# Path where we will save our training logs\n",
    "training_top_dir = Path(f\"{os.getcwd()}/data/\")\n",
    "# Create top_training_dir directory if needed, and launch tensorboard\n",
    "training_top_dir.mkdir(parents=True, exist_ok=True)\n",
    "log_dir = training_top_dir / \"06_image_translation/logs/\"\n",
    "# Create log directory if needed, and launch tensorboard\n",
    "log_dir.mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "if not data_path.exists():\n",
    "    raise FileNotFoundError(\n",
    "        f\"Data not found at {data_path}. Please check the top_dir and data_path variables.\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7e50558",
   "metadata": {
    "tags": []
   },
   "source": [
    "The next cell starts tensorboard.\n",
    "\n",
    "<div class=\"alert alert-warning\">\n",
    "If you launched jupyter lab from ssh terminal, add <code>--host &lt;your-server-name&gt;</code> to the tensorboard command below. <code>&lt;your-server-name&gt;</code> is the address of your compute node that ends in amazonaws.com.\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30625f14",
   "metadata": {
    "lines_to_next_cell": 0,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Imports and paths\n",
    "# Function to find an available port\n",
    "def find_free_port():\n",
    "    import socket\n",
    "\n",
    "    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:\n",
    "        s.bind((\"\", 0))\n",
    "        return s.getsockname()[1]\n",
    "\n",
    "\n",
    "# Launch TensorBoard on the browser\n",
    "def launch_tensorboard(log_dir):\n",
    "    import subprocess\n",
    "\n",
    "    port = find_free_port()\n",
    "    tensorboard_cmd = f\"tensorboard --logdir={log_dir} --port={port}\"\n",
    "    process = subprocess.Popen(tensorboard_cmd, shell=True)\n",
    "    print(\n",
    "        f\"TensorBoard started at http://localhost:{port}. \\n\"\n",
    "        \"If you are using VSCode remote session, forward the port using the PORTS tab next to TERMINAL.\"\n",
    "    )\n",
    "    return process\n",
    "\n",
    "\n",
    "# Launch tensorboard and click on the link to view the logs.\n",
    "tensorboard_process = launch_tensorboard(log_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c608b84",
   "metadata": {
    "lines_to_next_cell": 2,
    "tags": []
   },
   "source": [
    "<div class=\"alert alert-warning\">\n",
    "If you are using VSCode and a remote server, you will need to forward the port to view the tensorboard. <br>\n",
    "Take note of the port number was assigned in the previous cell.(i.e <code> http://localhost:{port_number_assigned}</code>) <br>\n",
    "\n",
    "Locate the your VSCode terminal and select the <code>Ports</code> tab <br>\n",
    "<ul>\n",
    "<li>Add a new port with the <code>port_number_assigned</code>\n",
    "</ul>\n",
    "Click on the link to view the tensorboard and it should open in your browser.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21fcca38",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Load OME-Zarr Dataset\n",
    "\n",
    "There should be 34 FOVs in the dataset.\n",
    "\n",
    "Each FOV consists of 3 channels of 2048x2048 images,\n",
    "saved in the [High-Content Screening (HCS) layout](https://ngff.openmicroscopy.org/latest/#hcs-layout)\n",
    "specified by the Open Microscopy Environment Next Generation File Format\n",
    "(OME-NGFF).\n",
    "\n",
    "The 3 channels correspond to the QPI, nuclei, and cell membrane. The nuclei were stained with DAPI and the cell membrane with Cellmask.\n",
    "\n",
    "- The layout on the disk is: `row/col/field/pyramid_level/timepoint/channel/z/y/x.`\n",
    "- These datasets only have 1 level in the pyramid (highest resolution) which is '0'."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4069ff37",
   "metadata": {
    "lines_to_next_cell": 0,
    "tags": []
   },
   "source": [
    "<div class=\"alert alert-warning\">\n",
    "You can inspect the tree structure by using your terminal:\n",
    "<code> iohub info -v \"path-to-ome-zarr\" </code>\n",
    "\n",
    "<br>\n",
    "More info on the CLI:\n",
    "<code>iohub info --help </code> to see the help menu.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94db2e35",
   "metadata": {},
   "outputs": [],
   "source": [
    "# This is the python function called by `iohub info` CLI command\n",
    "print_info(data_path, verbose=True)\n",
    "\n",
    "# Open and inspect the dataset.\n",
    "dataset = open_ome_zarr(data_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3cbbf67",
   "metadata": {
    "lines_to_next_cell": 0,
    "tags": []
   },
   "source": [
    "<div class=\"alert alert-info\">\n",
    "\n",
    "### Task 1.1\n",
    "Look at a couple different fields of view (FOVs) by changing the `field` variable.\n",
    "Check the cell density, the cell morphologies, and fluorescence signal.\n",
    "HINT: look at the HCS Plate format to see what are your options.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c762f267",
   "metadata": {
    "lines_to_next_cell": 1,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Use the field and pyramid_level below to visualize data.\n",
    "row = 0\n",
    "col = 0\n",
    "field = 9  # TODO: Change this to explore data.\n",
    "\n",
    "# NOTE: this dataset only has one level\n",
    "pyaramid_level = 0\n",
    "\n",
    "# `channel_names` is the metadata that is stored with data according to the OME-NGFF spec.\n",
    "n_channels = len(dataset.channel_names)\n",
    "\n",
    "image = dataset[f\"{row}/{col}/{field}/{pyaramid_level}\"].numpy()\n",
    "print(f\"data shape: {image.shape}, FOV: {field}, pyramid level: {pyaramid_level}\")\n",
    "\n",
    "figure, axes = plt.subplots(1, n_channels, figsize=(9, 3))\n",
    "\n",
    "for i in range(n_channels):\n",
    "    for i in range(n_channels):\n",
    "        channel_image = image[0, i, 0]\n",
    "        # Adjust contrast to 0.5th and 99.5th percentile of pixel values.\n",
    "        p_low, p_high = np.percentile(channel_image, (0.5, 99.5))\n",
    "        channel_image = np.clip(channel_image, p_low, p_high)\n",
    "        axes[i].imshow(channel_image, cmap=\"gray\")\n",
    "        axes[i].axis(\"off\")\n",
    "        axes[i].set_title(dataset.channel_names[i])\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43cb5bef",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Explore the effects of augmentation on batch.\n",
    "\n",
    "VisCy builds on top of PyTorch Lightning. PyTorch Lightning is a thin wrapper around PyTorch that allows rapid experimentation. It provides a [DataModule](https://lightning.ai/docs/pytorch/stable/data/datamodule.html) to handle loading and processing of data during training. VisCy provides a child class, `HCSDataModule` to make it intuitve to access data stored in the HCS layout.\n",
    "\n",
    "The dataloader in `HCSDataModule` returns a batch of samples. A `batch` is a list of dictionaries. The length of the list is equal to the batch size. Each dictionary consists of following key-value pairs.\n",
    "- `source`: the input image, a tensor of size `(1, 1, Y, X)`\n",
    "- `target`: the target image, a tensor of size `(2, 1, Y, X)`\n",
    "- `index` : the tuple of (location of field in HCS layout, time, and z-slice) of the sample."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "93f5109d",
   "metadata": {
    "tags": []
   },
   "source": [
    "<div class=\"alert alert-info\">\n",
    "\n",
    "### Task 1.2\n",
    "- Run the next cell to setup a logger for your augmentations.\n",
    "- Setup the `HCSDataloader()` in for training.\n",
    "  - Configure the dataloader for the `\"UNeXt2_2D\"`\n",
    "  - Configure the dataloader for the phase (source) to fluorescence cell nuclei and membrane (targets) regression task.\n",
    "  - Configure the dataloader for training. Hint: use the `HCSDataloader.setup()`\n",
    "- Open your tensorboard and look at the `IMAGES tab`.\n",
    "\n",
    "Note: If tensorboard is not showing images or the plots, try refreshing and using the \"Images\" tab.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cfba661d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a function to write a batch to tensorboard log.\n",
    "def log_batch_tensorboard(batch, batchno, writer, card_name):\n",
    "    \"\"\"\n",
    "    Logs a batch of images to TensorBoard.\n",
    "\n",
    "    Args:\n",
    "        batch (dict): A dictionary containing the batch of images to be logged.\n",
    "        writer (SummaryWriter): A TensorBoard SummaryWriter object.\n",
    "        card_name (str): The name of the card to be displayed in TensorBoard.\n",
    "\n",
    "    Returns:\n",
    "        None\n",
    "    \"\"\"\n",
    "    batch_phase = batch[\"source\"][:, :, 0, :, :]  # batch_size x z_size x Y x X tensor.\n",
    "    batch_membrane = batch[\"target\"][:, 1, 0, :, :].unsqueeze(\n",
    "        1\n",
    "    )  # batch_size x 1 x Y x X tensor.\n",
    "    batch_nuclei = batch[\"target\"][:, 0, 0, :, :].unsqueeze(\n",
    "        1\n",
    "    )  # batch_size x 1 x Y x X tensor.\n",
    "\n",
    "    p1, p99 = np.percentile(batch_membrane, (0.1, 99.9))\n",
    "    batch_membrane = np.clip((batch_membrane - p1) / (p99 - p1), 0, 1)\n",
    "\n",
    "    p1, p99 = np.percentile(batch_nuclei, (0.1, 99.9))\n",
    "    batch_nuclei = np.clip((batch_nuclei - p1) / (p99 - p1), 0, 1)\n",
    "\n",
    "    p1, p99 = np.percentile(batch_phase, (0.1, 99.9))\n",
    "    batch_phase = np.clip((batch_phase - p1) / (p99 - p1), 0, 1)\n",
    "\n",
    "    [N, C, H, W] = batch_phase.shape\n",
    "    interleaved_images = torch.zeros((3 * N, C, H, W), dtype=batch_phase.dtype)\n",
    "    interleaved_images[0::3, :] = batch_phase\n",
    "    interleaved_images[1::3, :] = batch_nuclei\n",
    "    interleaved_images[2::3, :] = batch_membrane\n",
    "\n",
    "    grid = torchvision.utils.make_grid(interleaved_images, nrow=3)\n",
    "\n",
    "    # add the grid to tensorboard\n",
    "    writer.add_image(card_name, grid, batchno)\n",
    "\n",
    "# Define a function to visualize a batch on jupyter, in case tensorboard is finicky\n",
    "def log_batch_jupyter(batch):\n",
    "    \"\"\"\n",
    "    Logs a batch of images on jupyter using ipywidget.\n",
    "\n",
    "    Args:\n",
    "        batch (dict): A dictionary containing the batch of images to be logged.\n",
    "\n",
    "    Returns:\n",
    "        None\n",
    "    \"\"\"\n",
    "    batch_phase = batch[\"source\"][:, :, 0, :, :]  # batch_size x z_size x Y x X tensor.\n",
    "    batch_size = batch_phase.shape[0]\n",
    "    batch_membrane = batch[\"target\"][:, 1, 0, :, :].unsqueeze(\n",
    "        1\n",
    "    )  # batch_size x 1 x Y x X tensor.\n",
    "    batch_nuclei = batch[\"target\"][:, 0, 0, :, :].unsqueeze(\n",
    "        1\n",
    "    )  # batch_size x 1 x Y x X tensor.\n",
    "\n",
    "    p1, p99 = np.percentile(batch_membrane, (0.1, 99.9))\n",
    "    batch_membrane = np.clip((batch_membrane - p1) / (p99 - p1), 0, 1)\n",
    "\n",
    "    p1, p99 = np.percentile(batch_nuclei, (0.1, 99.9))\n",
    "    batch_nuclei = np.clip((batch_nuclei - p1) / (p99 - p1), 0, 1)\n",
    "\n",
    "    p1, p99 = np.percentile(batch_phase, (0.1, 99.9))\n",
    "    batch_phase = np.clip((batch_phase - p1) / (p99 - p1), 0, 1)\n",
    "\n",
    "    n_channels = batch[\"target\"].shape[1] + batch[\"source\"].shape[1]\n",
    "    plt.figure()\n",
    "    fig, axes = plt.subplots(\n",
    "        batch_size, n_channels, figsize=(n_channels * 2, batch_size * 2)\n",
    "    )\n",
    "    [N, C, H, W] = batch_phase.shape\n",
    "    for sample_id in range(batch_size):\n",
    "        axes[sample_id, 0].imshow(batch_phase[sample_id, 0])\n",
    "        axes[sample_id, 1].imshow(batch_nuclei[sample_id, 0])\n",
    "        axes[sample_id, 2].imshow(batch_membrane[sample_id, 0])\n",
    "\n",
    "        for i in range(n_channels):\n",
    "            axes[sample_id, i].axis(\"off\")\n",
    "            axes[sample_id, i].set_title(dataset.channel_names[i])\n",
    "    plt.tight_layout()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2878a40d",
   "metadata": {
    "lines_to_next_cell": 0,
    "tags": [
     "solution"
    ]
   },
   "outputs": [],
   "source": [
    "# #######################\n",
    "# ##### SOLUTION ########\n",
    "# #######################\n",
    "\n",
    "BATCH_SIZE = 4\n",
    "# 4 is a perfectly reasonable batch size\n",
    "# (batch size does not have to be a power of 2)\n",
    "# See: https://sebastianraschka.com/blog/2022/batch-size-2.html\n",
    "\n",
    "source_channel = [\"Phase3D\"]\n",
    "target_channel = [\"Nucl\", \"Mem\"]\n",
    "\n",
    "data_module = HCSDataModule(\n",
    "    data_path,\n",
    "    z_window_size=1,\n",
    "    architecture=\"UNeXt2_2D\",\n",
    "    source_channel=source_channel,\n",
    "    target_channel=target_channel,\n",
    "    split_ratio=0.8,\n",
    "    batch_size=BATCH_SIZE,\n",
    "    num_workers=8,\n",
    "    yx_patch_size=(256, 256),  # larger patch size makes it easy to see augmentations.\n",
    "    augmentations=[],  # Turn off augmentation for now.\n",
    "    normalizations=[],  # Turn off normalization for now.\n",
    ")\n",
    "\n",
    "# Setup the data_module to fit. HINT: data_module.setup()\n",
    "data_module.setup(\"fit\")\n",
    "\n",
    "# Evaluate the data module\n",
    "print(\n",
    "    f\"Samples in training set: {len(data_module.train_dataset)}, \"\n",
    "    f\"samples in validation set:{len(data_module.val_dataset)}\"\n",
    ")\n",
    "train_dataloader = data_module.train_dataloader()\n",
    "# Instantiate the tensorboard SummaryWriter, logs the first batch and then iterates through all the batches and logs them to tensorboard.\n",
    "writer = SummaryWriter(log_dir=f\"{log_dir}/view_batch\")\n",
    "# Draw a batch and write to tensorboard.\n",
    "batch = next(iter(train_dataloader))\n",
    "log_batch_tensorboard(batch, 0, writer, \"augmentation/none\")\n",
    "writer.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c9884f12",
   "metadata": {
    "tags": []
   },
   "source": [
    "<div class=\"alert alert-warning\">\n",
    "\n",
    "### Questions\n",
    "1. What are the two channels in the target image?\n",
    "2. How many samples are in the training and validation set? What determined that split?\n",
    "\n",
    "Note: If tensorboard is not showing images, try refreshing and using the \"Images\" tab.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "376179a4",
   "metadata": {
    "lines_to_next_cell": 0,
    "tags": []
   },
   "source": [
    "If your tensorboard is causing issues, you can visualize directly on Jupyter /VSCode"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6685420f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize in Jupyter\n",
    "log_batch_jupyter(batch)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "03e19e4d",
   "metadata": {
    "lines_to_next_cell": 0,
    "tags": []
   },
   "source": [
    "<div class=\"alert alert-warning\">\n",
    "<h3> Question for Task 1.3 </h3>\n",
    "1. How do they make the model more robust to imaging parameters or conditions\n",
    "without having to acquire data for every possible condition? <br>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca2ca467",
   "metadata": {
    "lines_to_next_cell": 0,
    "tags": []
   },
   "source": [
    "<div class=\"alert alert-info\">\n",
    "\n",
    "### Task 1.3\n",
    "Add the following augmentations: \n",
    "- Add augmentations to rotate about $\\pi$ around z-axis, 30% scale in y,x,\n",
    "shearing of 10% and no padding with zeros with a probablity of 80%.\n",
    "- Add a Gaussian noise with a mean of 0.0 and standard deviation of 0.3 with a probability of 50%.\n",
    "\n",
    "HINT: `RandAffined()` and `RandGaussianNoised()` are from\n",
    "`viscy.transforms` [here](https://github.com/mehta-lab/VisCy/blob/main/viscy/transforms.py). You can look at the docs by running `RandAffined?`.<br><br>\n",
    "*Note these are MONAI transforms that have been redefined for VisCy.* \n",
    "[Compare your choice of augmentations by dowloading the pretrained models and config files](https://github.com/mehta-lab/VisCy/releases/download/v0.1.0/VisCy-0.1.0-VS-models.zip).\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5209d705",
   "metadata": {
    "tags": [
     "solution"
    ]
   },
   "outputs": [],
   "source": [
    "# #######################\n",
    "# ##### SOLUTION ########\n",
    "# #######################\n",
    "source_channel = [\"Phase3D\"]\n",
    "target_channel = [\"Nucl\", \"Mem\"]\n",
    "\n",
    "augmentations = [\n",
    "    RandWeightedCropd(\n",
    "        keys=source_channel + target_channel,\n",
    "        spatial_size=(1, 384, 384),\n",
    "        num_samples=2,\n",
    "        w_key=target_channel[0],\n",
    "    ),\n",
    "    RandAffined(\n",
    "        keys=source_channel + target_channel,\n",
    "        rotate_range=[3.14, 0.0, 0.0],\n",
    "        scale_range=[0.0, 0.3, 0.3],\n",
    "        prob=0.8,\n",
    "        padding_mode=\"zeros\",\n",
    "        shear_range=[0.0, 0.01, 0.01],\n",
    "    ),\n",
    "    RandAdjustContrastd(keys=source_channel, prob=0.5, gamma=(0.8, 1.2)),\n",
    "    RandScaleIntensityd(keys=source_channel, factors=0.5, prob=0.5),\n",
    "    RandGaussianNoised(keys=source_channel, prob=0.5, mean=0.0, std=0.3),\n",
    "    RandGaussianSmoothd(\n",
    "        keys=source_channel,\n",
    "        sigma_x=(0.25, 0.75),\n",
    "        sigma_y=(0.25, 0.75),\n",
    "        sigma_z=(0.0, 0.0),\n",
    "        prob=0.5,\n",
    "    ),\n",
    "]\n",
    "\n",
    "normalizations = [\n",
    "    NormalizeSampled(\n",
    "        keys=source_channel,\n",
    "        level=\"fov_statistics\",\n",
    "        subtrahend=\"mean\",\n",
    "        divisor=\"std\",\n",
    "    ),\n",
    "    NormalizeSampled(\n",
    "        keys=target_channel,\n",
    "        level=\"fov_statistics\",\n",
    "        subtrahend=\"median\",\n",
    "        divisor=\"iqr\",\n",
    "    )\n",
    "]\n",
    "\n",
    "data_module.augmentations = augmentations\n",
    "\n",
    "# Setup the data_module to fit. HINT: data_module.setup()\n",
    "data_module.setup(\"fit\")\n",
    "\n",
    "# get the new data loader with augmentation turned on\n",
    "augmented_train_dataloader = data_module.train_dataloader()\n",
    "\n",
    "# Draw batches and write to tensorboard\n",
    "writer = SummaryWriter(log_dir=f\"{log_dir}/view_batch\")\n",
    "augmented_batch = next(iter(augmented_train_dataloader))\n",
    "log_batch_tensorboard(augmented_batch, 0, writer, \"augmentation/some\")\n",
    "writer.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1be24b3",
   "metadata": {
    "tags": []
   },
   "source": [
    "<div class=\"alert alert-warning\">\n",
    "<h3> Question for Task 1.3 </h3>\n",
    "1. Look at your tensorboard. Can you tell the agumentations were applied to the sample batch? Compare the batch with and without augmentations. <br>\n",
    "2. Are these augmentations good enough? What else would you add?\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dcdb1d7c",
   "metadata": {},
   "source": [
    "Visualize directly on Jupyter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fedf6687",
   "metadata": {},
   "outputs": [],
   "source": [
    "log_batch_jupyter(augmented_batch)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0cd143e",
   "metadata": {
    "lines_to_next_cell": 0,
    "tags": []
   },
   "source": [
    "## Train a 2D U-Net model to predict nuclei and membrane from phase.\n",
    "### Constructing a 2D UNeXt2 using VisCy"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a4f90e03",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-info\">\n",
    "\n",
    "### Task 1.5\n",
    "- Run the next cell to instantiate the `UNeXt2_2D` model\n",
    "  - Configure the network for the phase (source) to fluorescence cell nuclei and membrane (targets) regression task.\n",
    "  - Call the VSUNet with the `\"UNeXt2_2D\"` architecture.\n",
    "- Run the next cells to instantiate data module and trainer.\n",
    "  - Add the source channel name and the target channel names\n",
    "- Start the training <br>\n",
    "\n",
    "<b> Note </b> <br>\n",
    "See ``viscy.unet.networks.Unet2D.Unet2d`` ([source code](https://github.com/mehta-lab/VisCy/blob/7c5e4c1d68e70163cf514d22c475da8ea7dc3a88/viscy/unet/networks/Unet2D.py#L7)) to learn more about the configuration.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a589abac",
   "metadata": {
    "tags": [
     "solution"
    ]
   },
   "outputs": [],
   "source": [
    "\n",
    "# Here we are creating a 2D UNet.\n",
    "GPU_ID = 0\n",
    "\n",
    "BATCH_SIZE = 16\n",
    "YX_PATCH_SIZE = (256, 256)\n",
    "\n",
    "# Dictionary that specifies key parameters of the model.\n",
    "# #######################\n",
    "# ##### SOLUTION ########\n",
    "# #######################\n",
    "phase2fluor_config = dict(\n",
    "    in_channels=1,\n",
    "    out_channels=2,\n",
    "    encoder_blocks=[3, 3, 9, 3],\n",
    "    dims=[96, 192, 384, 768],\n",
    "    decoder_conv_blocks=2,\n",
    "    stem_kernel_size=(1, 2, 2),\n",
    "    in_stack_depth=1,\n",
    "    pretraining=False,\n",
    ")\n",
    "\n",
    "phase2fluor_model = VSUNet(\n",
    "    architecture=\"UNeXt2_2D\",  # 2D UNeXt2 architecture\n",
    "    model_config=phase2fluor_config.copy(),\n",
    "    loss_function=MixedLoss(l1_alpha=0.5, l2_alpha=0.0, ms_dssim_alpha=0.5),\n",
    "    schedule=\"WarmupCosine\",\n",
    "    lr=6e-4,\n",
    "    log_batches_per_epoch=5,  # Number of samples from each batch to log to tensorboard.\n",
    "    freeze_encoder=False,\n",
    ")\n",
    "\n",
    "# ### Instantiate data module and trainer, test that we are setup to launch training.\n",
    "# #######################\n",
    "# ##### SOLUTION ########\n",
    "# #######################\n",
    "# Selecting the source and target channel names from the dataset.\n",
    "source_channel = [\"Phase3D\"]\n",
    "target_channel = [\"Nucl\", \"Mem\"]\n",
    "# Setup the data module.\n",
    "phase2fluor_2D_data = HCSDataModule(\n",
    "    data_path,\n",
    "    architecture=\"UNeXt2_2D\",\n",
    "    source_channel=source_channel,\n",
    "    target_channel=target_channel,\n",
    "    z_window_size=1,\n",
    "    split_ratio=0.8,\n",
    "    batch_size=BATCH_SIZE,\n",
    "    num_workers=8,\n",
    "    yx_patch_size=YX_PATCH_SIZE,\n",
    "    augmentations=augmentations,\n",
    "    normalizations=normalizations,\n",
    ")\n",
    "# #######################\n",
    "# ##### SOLUTION ########\n",
    "# #######################\n",
    "phase2fluor_2D_data.setup(\"fit\")\n",
    "# fast_dev_run runs a single batch of data through the model to check for errors.\n",
    "trainer = VSTrainer(accelerator=\"gpu\", devices=[GPU_ID],precision='16-mixed', fast_dev_run=True)\n",
    "\n",
    "# trainer class takes the model and the data module as inputs.\n",
    "trainer.fit(phase2fluor_model, datamodule=phase2fluor_2D_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "968ab08a",
   "metadata": {
    "tags": []
   },
   "source": [
    "## View model graph.\n",
    "\n",
    "PyTorch uses dynamic graphs under the hood.\n",
    "The graphs are constructed on the fly.\n",
    "This is in contrast to TensorFlow,\n",
    "where the graph is constructed before the training loop and remains static.\n",
    "In other words, the graph of the network can change with every forward pass.\n",
    "Therefore, we need to supply an input tensor to construct the graph.\n",
    "The input tensor can be a random tensor of the correct shape and type.\n",
    "We can also supply a real image from the dataset.\n",
    "The latter is more useful for debugging."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb904146",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-info\">\n",
    "\n",
    "### Task 1.5\n",
    "Run the next cell to generate a graph representation of the model architecture.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43ab6504",
   "metadata": {},
   "outputs": [],
   "source": [
    "# visualize graph of phase2fluor model as image.\n",
    "model_graph_phase2fluor = torchview.draw_graph(\n",
    "    phase2fluor_model,\n",
    "    phase2fluor_2D_data.train_dataset[0][\"source\"][0].unsqueeze(dim=0),\n",
    "    roll=True,\n",
    "    depth=3,  # adjust depth to zoom in.\n",
    "    device=\"cpu\",\n",
    "    # expand_nested=True,\n",
    ")\n",
    "# Print the image of the model.\n",
    "model_graph_phase2fluor.visual_graph"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c468747",
   "metadata": {
    "tags": []
   },
   "source": [
    "<div class=\"alert alert-warning\">\n",
    "\n",
    "### Question:\n",
    "Can you recognize the UNet structure and skip connections in this graph visualization?\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a38efb8",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-info\">\n",
    "\n",
    "<h3> Task 1.6 </h3>\n",
    "Start training by running the following cell. Check the new logs on the tensorboard.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b5539a48",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [],
   "source": [
    "# Check if GPU is available\n",
    "# You can check by typing `nvidia-smi`\n",
    "GPU_ID = 0\n",
    "\n",
    "n_samples = len(phase2fluor_2D_data.train_dataset)\n",
    "steps_per_epoch = n_samples // BATCH_SIZE  # steps per epoch.\n",
    "n_epochs = 80  # Set this to 80-100 or the number of epochs you want to train for.\n",
    "\n",
    "trainer = VSTrainer(\n",
    "    accelerator=\"gpu\",\n",
    "    devices=[GPU_ID],\n",
    "    max_epochs=n_epochs,\n",
    "    precision='16-mixed',\n",
    "    log_every_n_steps=steps_per_epoch // 2,\n",
    "    # log losses and image samples 2 times per epoch.\n",
    "    logger=TensorBoardLogger(\n",
    "        save_dir=log_dir,\n",
    "        # lightning trainer transparently saves logs and model checkpoints in this directory.\n",
    "        name=\"phase2fluor\",\n",
    "        log_graph=True,\n",
    "    ),\n",
    ")\n",
    "# Launch training and check that loss and images are being logged on tensorboard.\n",
    "trainer.fit(phase2fluor_model, datamodule=phase2fluor_2D_data)\n",
    "\n",
    "# Move the model to the GPU.\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "phase2fluor_model.to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e9e5243",
   "metadata": {
    "lines_to_next_cell": 0,
    "tags": []
   },
   "source": [
    "<div class=\"alert alert-success\">\n",
    "\n",
    "<h2> Checkpoint 1 </h2>\n",
    "\n",
    "While your model is training, let's think about the following questions:<br>\n",
    "<ul>\n",
    "<li>What is the information content of each channel in the dataset?</li>\n",
    "<li>How would you use image translation models?</li>\n",
    "<li>What can you try to improve the performance of each model?</li>\n",
    "</ul>\n",
    "\n",
    "Now the training has started,\n",
    "we can come back after a while and evaluate the performance!\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e2043d4",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Part 2: Assess your trained model\n",
    "\n",
    "Now we will look at some metrics of performance of previous model.\n",
    "We typically evaluate the model performance on a held out test data.\n",
    "We will use the following metrics to evaluate the accuracy of regression of the model:\n",
    "\n",
    "- [Person Correlation](https://en.wikipedia.org/wiki/Pearson_correlation_coefficient).\n",
    "- [Structural similarity](https://en.wikipedia.org/wiki/Structural_similarity) (SSIM). \n",
    "\n",
    "You should also look at the validation samples on tensorboard\n",
    "(hint: the experimental data in nuclei channel is imperfect.)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3263b1e8",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-info\">\n",
    "\n",
    "<h3> Task 2.1 Define metrics </h3>\n",
    "\n",
    "For each of the above metrics, write a brief definition of what they are and what they mean\n",
    "for this image translation task. Use your favorite search engine and/or resources.\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d71ce538",
   "metadata": {
    "tags": []
   },
   "source": [
    "```\n",
    "#######################\n",
    "##### Todo ############\n",
    "#######################\n",
    "\n",
    "```\n",
    "\n",
    "- Pearson Correlation:\n",
    "\n",
    "- Structural similarity:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d3fd630",
   "metadata": {
    "lines_to_next_cell": 0,
    "tags": []
   },
   "source": [
    "### Let's compute metrics directly and plot below."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c643c8bd",
   "metadata": {
    "lines_to_next_cell": 0,
    "tags": []
   },
   "source": [
    "<div class=\"alert alert-danger\">\n",
    "If you weren't able to train or training didn't complete please run the following lines to load the latest checkpoint <br>\n",
    "\n",
    "```python\n",
    "phase2fluor_model_ckpt = natsorted(glob(\n",
    "   str(top_dir / \"06_image_translation/logs/phase2fluor/version*/checkpoints/*.ckpt\")\n",
    "))[-1]\n",
    "```\n",
    "<br>\n",
    "NOTE: if their model didn't go past epoch 5, lost their checkpoint, or didnt train anything. \n",
    "Run the following:\n",
    "\n",
    "```python\n",
    "phase2fluor_model_ckpt = natsorted(glob(\n",
    " str(top_dir/\"06_image_translation/backup/phase2fluor/version_0/checkpoints/*.ckpt\")\n",
    "))[-1]\n",
    "```\n",
    "\n",
    "```python\n",
    "phase2fluor_config = dict(\n",
    "    in_channels=1,\n",
    "    out_channels=2,\n",
    "    encoder_blocks=[3, 3, 9, 3],\n",
    "    dims=[96, 192, 384, 768],\n",
    "    decoder_conv_blocks=2,\n",
    "    stem_kernel_size=(1, 2, 2),\n",
    "    in_stack_depth=1,\n",
    "    pretraining=False,\n",
    ")\n",
    "Load the model checkpoint\n",
    "phase2fluor_model = VSUNet.load_from_checkpoint(\n",
    "    phase2fluor_model_ckpt,\n",
    "    architecture=\"UNeXt2_2D\",\n",
    "    model_config = phase2fluor_config,\n",
    "    accelerator='gpu'\n",
    ")\n",
    "````\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e53e4295",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "# Setup the test data module.\n",
    "test_data_path = top_dir / \"06_image_translation/test/a549_hoechst_cellmask_test.zarr\"\n",
    "source_channel = [\"Phase3D\"]\n",
    "target_channel = [\"Nucl\", \"Mem\"]\n",
    "\n",
    "test_data = HCSDataModule(\n",
    "    test_data_path,\n",
    "    source_channel=source_channel,\n",
    "    target_channel=target_channel,\n",
    "    z_window_size=1,\n",
    "    batch_size=1,\n",
    "    num_workers=8,\n",
    "    architecture=\"UNeXt2\",\n",
    ")\n",
    "test_data.setup(\"test\")\n",
    "\n",
    "test_metrics = pd.DataFrame(\n",
    "    columns=[\"pearson_nuc\", \"SSIM_nuc\", \"pearson_mem\", \"SSIM_mem\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0822d689",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "# Compute metrics directly and plot here.\n",
    "def normalize_fov(input:ArrayLike):\n",
    "    \"Normalizing the fov with zero mean and unit variance\"\n",
    "    mean = np.mean(input)\n",
    "    std = np.std(input)\n",
    "    return (input - mean) / std\n",
    "\n",
    "for i, sample in enumerate(tqdm(test_data.test_dataloader(), desc=\"Computing metrics per sample\")):\n",
    "    phase_image = sample[\"source\"].to(phase2fluor_model.device)\n",
    "    with torch.inference_mode():  # turn off gradient computation.\n",
    "        predicted_image = phase2fluor_model(phase_image)\n",
    "\n",
    "    target_image = (\n",
    "        sample[\"target\"].cpu().numpy().squeeze(0)\n",
    "    )  # Squeezing batch dimension.\n",
    "    predicted_image = predicted_image.cpu().numpy().squeeze(0)\n",
    "    phase_image = phase_image.cpu().numpy().squeeze(0)\n",
    "    target_mem = normalize_fov(target_image[1, 0, :, :])\n",
    "    target_nuc = normalize_fov(target_image[0, 0, :, :])\n",
    "    # slicing channel dimension, squeezing z-dimension.\n",
    "    predicted_mem = normalize_fov(predicted_image[1, :, :, :].squeeze(0))\n",
    "    predicted_nuc = normalize_fov(predicted_image[0, :, :, :].squeeze(0))\n",
    "\n",
    "    # Compute SSIM and pearson correlation.\n",
    "    ssim_nuc = metrics.structural_similarity(target_nuc, predicted_nuc, data_range=1)\n",
    "    ssim_mem = metrics.structural_similarity(target_mem, predicted_mem, data_range=1)\n",
    "    pearson_nuc = np.corrcoef(target_nuc.flatten(), predicted_nuc.flatten())[0, 1]\n",
    "    pearson_mem = np.corrcoef(target_mem.flatten(), predicted_mem.flatten())[0, 1]\n",
    "\n",
    "    test_metrics.loc[i] = {\n",
    "        \"pearson_nuc\": pearson_nuc,\n",
    "        \"SSIM_nuc\": ssim_nuc,\n",
    "        \"pearson_mem\": pearson_mem,\n",
    "        \"SSIM_mem\": ssim_mem,\n",
    "    }\n",
    "\n",
    "# Plot the following metrics\n",
    "test_metrics.boxplot(\n",
    "    column=[\"pearson_nuc\", \"SSIM_nuc\", \"pearson_mem\", \"SSIM_mem\"],\n",
    "    rot=30,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3aaba4e",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [],
   "source": [
    "# Adjust the image to the 0.5-99.5 percentile range.\n",
    "def process_image(image):\n",
    "    p_low, p_high = np.percentile(image, (0.5, 99.5))\n",
    "    return np.clip(image, p_low, p_high)\n",
    "\n",
    "# Plot the predicted image vs target image.\n",
    "channel_titles = [\"Phase\", \"Target Nuclei\", \"Target Membrane\", \"Predicted Nuclei\", \"Predicted Membrane\"]\n",
    "fig, axes = plt.subplots(5, 1, figsize=(20, 20))\n",
    "\n",
    "# Get a writer to output the images into tensorboard and plot the source, predictions and target images\n",
    "for i, sample in enumerate(test_data.test_dataloader()):\n",
    "    # Plot the phase image\n",
    "    phase_image = sample[\"source\"]\n",
    "    channel_image = phase_image[0, 0, 0]\n",
    "    p_low, p_high = np.percentile(channel_image, (0.5, 99.5))\n",
    "    channel_image = np.clip(channel_image, p_low, p_high)\n",
    "    axes[0].imshow(channel_image, cmap=\"gray\")\n",
    "    axes[0].axis(\"off\")\n",
    "    axes[0].set_title(channel_titles[0])\n",
    "\n",
    "    with torch.inference_mode():  # turn off gradient computation.\n",
    "        predicted_image = (\n",
    "            phase2fluor_model(phase_image.to(phase2fluor_model.device))\n",
    "            .cpu()\n",
    "            .numpy()\n",
    "            .squeeze(0)\n",
    "        )\n",
    "\n",
    "    target_image = sample[\"target\"].cpu().numpy().squeeze(0)\n",
    "    phase_raw = process_image(phase_image[0, 0, 0])\n",
    "    predicted_nuclei = process_image(predicted_image[0,0])\n",
    "    predicted_membrane = process_image(predicted_image[1,0])\n",
    "    target_nuclei = process_image(target_image[0,0])\n",
    "    target_membrane = process_image(target_image[1,0])\n",
    "       # Concatenate all images side by side\n",
    "    combined_image = np.concatenate(\n",
    "        (\n",
    "            phase_raw, \n",
    "            predicted_nuclei,\n",
    "            predicted_membrane,\n",
    "            target_nuclei,\n",
    "            target_membrane\n",
    "        ),\n",
    "        axis=1,\n",
    "    )\n",
    "\n",
    "    # Plot the phase,target nuclei, target membrane, predicted nuclei, predicted membrane\n",
    "    axes[1].imshow(target_nuclei, cmap=\"gray\")\n",
    "    axes[2].imshow(target_membrane, cmap=\"gray\")\n",
    "    axes[3].imshow(predicted_nuclei, cmap=\"gray\")\n",
    "    axes[4].imshow(predicted_membrane, cmap=\"gray\")\n",
    "\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "    break"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88faa26e",
   "metadata": {
    "lines_to_next_cell": 2,
    "tags": []
   },
   "source": [
    "<div class=\"alert alert-info\">\n",
    "\n",
    "<h3> Task 2.2 Loading the pretrained model VSCyto2D </h3>\n",
    "Here we will compare your model with the VSCyto2D pretrained model by computing the pixel-based metrics and segmentation-based metrics.\n",
    "\n",
    "<ul>\n",
    "<li>When you ran the `setup.sh` you also downloaded the models in `/06_image_translation/pretrained_models/VSCyto2D/*.ckpt`</li>\n",
    "<li>Load the <b>VSCyto2 model</b> model checkpoint and the configuration file</li>\n",
    "<li>Compute the pixel-based metrics and segmentation-based metrics between the model you trained and the pretrained model</li>\n",
    "</ul>\n",
    "<br>\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1c5193f",
   "metadata": {
    "lines_to_next_cell": 0,
    "tags": [
     "solution"
    ]
   },
   "outputs": [],
   "source": [
    "# #######################\n",
    "# ##### SOLUTION ########\n",
    "# #######################\n",
    "\n",
    "pretrained_model_ckpt = (\n",
    "    top_dir / \"06_image_translation/pretrained_models/VSCyto2D/epoch=399-step=23200.ckpt\"\n",
    ")\n",
    "\n",
    "phase2fluor_config = dict(\n",
    "    in_channels=1,\n",
    "    out_channels=2,\n",
    "    encoder_blocks=[3, 3, 9, 3],\n",
    "    dims=[96, 192, 384, 768],\n",
    "    decoder_conv_blocks=2,\n",
    "    stem_kernel_size=(1, 2, 2),\n",
    "    in_stack_depth=1,\n",
    "    pretraining=False,\n",
    ")\n",
    "# Load the model checkpoint\n",
    "pretrained_phase2fluor = VSUNet.load_from_checkpoint(\n",
    "    pretrained_model_ckpt,\n",
    "    architecture=\"UNeXt2_2D\",\n",
    "    model_config = phase2fluor_config,\n",
    "    accelerator='gpu'\n",
    ")\n",
    "pretrained_phase2fluor.eval()\n",
    "\n",
    "### Re-load your trained model\n",
    "# NOTE: assuming the latest checkpoint it your latest training and model\n",
    "phase2fluor_model_ckpt = natsorted(glob(\n",
    "    str(training_top_dir / \"06_image_translation/logs/phase2fluor/version*/checkpoints/*.ckpt\")\n",
    "))[-1]\n",
    "\n",
    "# NOTE: if their model didn't go past epoch 5, lost their checkpoint, or didnt train anything. \n",
    "# Uncomment the next lines\n",
    "#phase2fluor_model_ckpt = natsorted(glob(\n",
    "#  str(top_dir/\"06_image_translation/backup/phase2fluor/version_0/checkpoints/*.ckpt\")\n",
    "#))[-1]\n",
    "\n",
    "phase2fluor_config = dict(\n",
    "    in_channels=1,\n",
    "    out_channels=2,\n",
    "    encoder_blocks=[3, 3, 9, 3],\n",
    "    dims=[96, 192, 384, 768],\n",
    "    decoder_conv_blocks=2,\n",
    "    stem_kernel_size=(1, 2, 2),\n",
    "    in_stack_depth=1,\n",
    "    pretraining=False,\n",
    ")\n",
    "# Load the model checkpoint\n",
    "phase2fluor_model = VSUNet.load_from_checkpoint(\n",
    "    phase2fluor_model_ckpt,\n",
    "    architecture=\"UNeXt2_2D\",\n",
    "    model_config = phase2fluor_config,\n",
    "    accelerator='gpu'\n",
    ")\n",
    "phase2fluor_model.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2d58211",
   "metadata": {
    "lines_to_next_cell": 0,
    "tags": []
   },
   "source": [
    "<div class=\"alert alert-warning\">\n",
    "<h3> Question </h3> \n",
    "1. Can we evaluate a model's performance based on their segmentations?<br>\n",
    "2. Look up IoU or Jaccard index, dice coefficient, and AP metrics. LINK:https://metrics-reloaded.dkfz.de/metric-library <br>\n",
    "We will evaluate the performance of your trained model with a pre-trained model using pixel based metrics as above and\n",
    "segmantation based metrics including (mAP@0.5, dice, accuracy and jaccard index). <br>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38d8d215",
   "metadata": {
    "lines_to_next_cell": 0,
    "tags": [
     "solution"
    ]
   },
   "source": [
    "\n",
    "- <b> IoU (Intersection over Union): </b> Also referred to as the Jaccard index, is essentially a method to quantify the percent overlap between the target and predicted masks. \n",
    "It is calculated as the intersection of the target and predicted masks divided by the union of the target and predicted masks. <br>\n",
    "- <b> Dice Coefficient:</b> Metric used to evaluate the similarity between two sets.<br>\n",
    "It is calculated as twice the intersection of the target and predicted masks divided by the sum of the target and predicted masks.<br>\n",
    "- <b> mAP (mean Average Precision):</b>  The mean Average Precision (mAP) is a metric used to evaluate the performance of object detection models. \n",
    "It is calculated as the average precision across all classes and is used to measure the accuracy of the model in localizing objects.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8faf713d",
   "metadata": {
    "lines_to_next_cell": 0,
    "tags": []
   },
   "source": [
    "### Let's compute the metrics for the test dataset\n",
    "Before you run the following code, make sure you have the pretrained model loaded and the test data is ready.\n",
    "\n",
    "The following code will compute the following:\n",
    "- the pixel-based metrics  (pearson correlation, SSIM)\n",
    "- segmentation-based metrics (mAP@0.5, dice, accuracy, jaccard index)\n",
    "\n",
    "#### Note:\n",
    "- The segmentation-based metrics are computed using the cellpose stock `nuclei` model\n",
    "- The metrics will be store in the `test_pixel_metrics` and `test_segmentation_metrics` dataframes\n",
    "- The segmentations will be stored in the `segmentation_store` zarr file\n",
    "- Analyze the code while it runs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1bcc6683",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "# Define the function to compute the cellpose segmentation\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",
    "        \"channels\": [0, 0],\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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ef67524",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "# Setting the paths for the test data and the output segmentation\n",
    "test_data_path = top_dir / \"06_image_translation/test/a549_hoechst_cellmask_test.zarr\"\n",
    "output_segmentation_path= training_top_dir /\"06_image_translation/pretrained_model_segmentations.zarr\"\n",
    "\n",
    "# Creating the dataframes to store the pixel and segmentation metrics\n",
    "test_pixel_metrics = pd.DataFrame(\n",
    "    columns=[\"model\", \"fov\",\"pearson_nuc\", \"SSIM_nuc\", \"pearson_mem\", \"SSIM_mem\"]\n",
    ")\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",
    "# Opening the test dataset\n",
    "test_dataset = open_ome_zarr(test_data_path)\n",
    "\n",
    "# Creating an output store for the predictions and segmentations\n",
    "segmentation_store = open_ome_zarr(output_segmentation_path,channel_names=['nuc_pred','mem_pred','nuc_labels'],mode='w',layout='hcs')\n",
    "\n",
    "# Looking at the test dataset\n",
    "print('Test dataset:')\n",
    "test_dataset.print_tree()\n",
    "channel_names = test_dataset.channel_names\n",
    "print(f'Channel names: {channel_names}')\n",
    "\n",
    "# Finding the channel indices for the corresponding channel names\n",
    "phase_cidx = channel_names.index(\"Phase3D\")\n",
    "nuc_cidx = channel_names.index(\"Nucl\")\n",
    "mem_cidx =  channel_names.index(\"Mem\")\n",
    "nuc_label_cidx =  channel_names.index(\"nuclei_segmentation\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e9e5be5",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [],
   "source": [
    "def min_max_scale(image:ArrayLike)->ArrayLike:\n",
    "    \"Normalizing the image using min-max scaling\"\n",
    "    min_val = image.min()\n",
    "    max_val = image.max()\n",
    "    return (image - min_val) / (max_val - min_val)\n",
    "\n",
    "# Iterating through the test dataset positions to:\n",
    "positions = list(test_dataset.positions())\n",
    "total_positions = len(positions)\n",
    "\n",
    "# Initializing the progress bar with the total number of positions\n",
    "with tqdm(total=total_positions, desc=\"Processing FOVs\") as pbar:\n",
    "    # Iterating through the test dataset positions\n",
    "    for fov, pos in positions:\n",
    "        T,C,Z,Y,X = pos.data.shape\n",
    "        Z_slice = slice(Z//2,Z//2+1)\n",
    "        # Getting the arrays and the center slices\n",
    "        phase_image = pos.data[:,phase_cidx:phase_cidx+1,Z_slice]\n",
    "        target_nucleus =  pos.data[0,nuc_cidx:nuc_cidx+1,Z_slice]\n",
    "        target_membrane =  pos.data[0,mem_cidx:mem_cidx+1,Z_slice]\n",
    "        target_nuc_label = pos.data[0,nuc_label_cidx:nuc_label_cidx+1,Z_slice]\n",
    "\n",
    "        #normalize the phase\n",
    "        phase_image = normalize_fov(phase_image)\n",
    "        \n",
    "        # Running the prediction for both models\n",
    "        phase_image = torch.from_numpy(phase_image).type(torch.float32)\n",
    "        phase_image = phase_image.to(phase2fluor_model.device)\n",
    "        with torch.inference_mode():  # turn off gradient computation.\n",
    "            predicted_image_phase2fluor = phase2fluor_model(phase_image)\n",
    "            predicted_image_pretrained = pretrained_phase2fluor(phase_image)\n",
    "\n",
    "        # Loading and Normalizing the target and predictions for both models \n",
    "        predicted_image_phase2fluor = predicted_image_phase2fluor.cpu().numpy().squeeze(0)\n",
    "        predicted_image_pretrained = predicted_image_pretrained.cpu().numpy().squeeze(0)\n",
    "        phase_image = phase_image.cpu().numpy().squeeze(0)\n",
    "\n",
    "        target_mem = min_max_scale(target_membrane[0,0])\n",
    "        target_nuc = min_max_scale(target_nucleus[0,0])\n",
    "    \n",
    "        # Normalizing the dataset using min-max scaling\n",
    "        predicted_mem_phase2fluor = min_max_scale(\n",
    "            predicted_image_phase2fluor[1, :, :, :].squeeze(0)\n",
    "        )\n",
    "        predicted_nuc_phase2fluor = min_max_scale(\n",
    "            predicted_image_phase2fluor[0, :, :, :].squeeze(0)\n",
    "        )\n",
    "\n",
    "        predicted_mem_pretrained = min_max_scale(\n",
    "            predicted_image_pretrained[1, :, :, :].squeeze(0)\n",
    "        )\n",
    "        predicted_nuc_pretrained = min_max_scale(\n",
    "            predicted_image_pretrained[0, :, :, :].squeeze(0)\n",
    "        )\n",
    "\n",
    "        #######  Pixel-based Metrics ############\n",
    "        # Compute SSIM and Pearson correlation for phase2fluor_model\n",
    "        print('Computing Pixel Metrics')\n",
    "        ssim_nuc_phase2fluor = metrics.structural_similarity(\n",
    "            target_nuc, predicted_nuc_phase2fluor, data_range=1\n",
    "        )\n",
    "        ssim_mem_phase2fluor = metrics.structural_similarity(\n",
    "            target_mem, predicted_mem_phase2fluor, data_range=1\n",
    "        )\n",
    "        pearson_nuc_phase2fluor = np.corrcoef(\n",
    "            target_nuc.flatten(), predicted_nuc_phase2fluor.flatten()\n",
    "        )[0, 1]\n",
    "        pearson_mem_phase2fluor = np.corrcoef(\n",
    "            target_mem.flatten(), predicted_mem_phase2fluor.flatten()\n",
    "        )[0, 1]\n",
    "\n",
    "        test_pixel_metrics.loc[len(test_pixel_metrics)] = {\n",
    "            \"model\": \"phase2fluor\",\n",
    "            \"fov\":fov,\n",
    "            \"pearson_nuc\": pearson_nuc_phase2fluor,\n",
    "            \"SSIM_nuc\": ssim_nuc_phase2fluor,\n",
    "            \"pearson_mem\": pearson_mem_phase2fluor,\n",
    "            \"SSIM_mem\": ssim_mem_phase2fluor,\n",
    "        }\n",
    "        # Compute SSIM and Pearson correlation for pretrained_model\n",
    "        ssim_nuc_pretrained = metrics.structural_similarity(\n",
    "            target_nuc, predicted_nuc_pretrained, data_range=1\n",
    "        )\n",
    "        ssim_mem_pretrained = metrics.structural_similarity(\n",
    "            target_mem, predicted_mem_pretrained, data_range=1\n",
    "        )\n",
    "        pearson_nuc_pretrained = np.corrcoef(\n",
    "            target_nuc.flatten(), predicted_nuc_pretrained.flatten()\n",
    "        )[0, 1]\n",
    "        pearson_mem_pretrained = np.corrcoef(\n",
    "            target_mem.flatten(), predicted_mem_pretrained.flatten()\n",
    "        )[0, 1]\n",
    "\n",
    "        test_pixel_metrics.loc[len(test_pixel_metrics)] = {\n",
    "            \"model\": \"pretrained_phase2fluor\",\n",
    "            \"fov\":fov,\n",
    "            \"pearson_nuc\": pearson_nuc_pretrained,\n",
    "            \"SSIM_nuc\": ssim_nuc_pretrained,\n",
    "            \"pearson_mem\": pearson_mem_pretrained,\n",
    "            \"SSIM_mem\": ssim_mem_pretrained,\n",
    "        }\n",
    "\n",
    "        ###### Segmentation based metrics #########\n",
    "        # Load the manually curated nuclei target label\n",
    "        print('Computing Segmentation Metrics')\n",
    "        pred_label,target_label= cellpose_segmentation(predicted_nuc_phase2fluor,target_nucleus)\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",
    "\n",
    "        test_segmentation_metrics.loc[len(test_segmentation_metrics)] = {\n",
    "            \"model\": \"phase2fluor\",\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",
    "\n",
    "        pred_label,target_label= cellpose_segmentation(predicted_nuc_pretrained,target_nucleus)\n",
    "        \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",
    "\n",
    "        test_segmentation_metrics.loc[len(test_segmentation_metrics)] = {\n",
    "            \"model\": \"phase2fluor_pretrained\",\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",
    "        \n",
    "        #Save the predictions and segmentations\n",
    "        position = segmentation_store.create_position(*Path(fov).parts[-3:])\n",
    "        output_array = np.zeros((T,3,1,Y,X),dtype=np.float32)\n",
    "        output_array[0,0,0]=predicted_nuc_pretrained\n",
    "        output_array[0,1,0]=predicted_mem_pretrained\n",
    "        output_array[0,2,0]=np.array(pred_label)\n",
    "        position.create_image(\"0\",output_array)\n",
    "\n",
    "        # Update the progress bar\n",
    "        pbar.update(1)\n",
    "    \n",
    "# Close the OME-Zarr files\n",
    "test_dataset.close()\n",
    "segmentation_store.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94fc466d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save the test metrics into a dataframe\n",
    "pixel_metrics_path = training_top_dir/\"06_image_translation/VS_metrics_pixel.csv\"\n",
    "segmentation_metrics_path = training_top_dir/\"06_image_translation/VS_metrics_segments.csv\"\n",
    "test_pixel_metrics.to_csv(pixel_metrics_path)\n",
    "test_segmentation_metrics.to_csv(segmentation_metrics_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0cdf8e08",
   "metadata": {
    "tags": []
   },
   "source": [
    "<div class=\"alert alert-info\">\n",
    "\n",
    "<h3> Task 2.3 Compare the model's metrics </h3>\n",
    "In the previous section, we computed the pixel-based metrics and segmentation-based metrics.\n",
    "Now we will compare the performance of the model you trained with the pretrained model by plotting the boxplots.\n",
    "\n",
    "After you plot the metrics answer the following:\n",
    "<ul>\n",
    "<li>What do these metrics tells us about the performance of the model?</li>\n",
    "<li>How do you interpret the differences in the metrics between the models?</li>\n",
    "<li>How is your model compared to the pretrained model? How can you improve it?</li>\n",
    "</ul>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b16b6b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Show boxplot of the metrics\n",
    "# Boxplot of the metrics\n",
    "test_pixel_metrics.boxplot(\n",
    "    by=\"model\",\n",
    "    column=[\"pearson_nuc\", \"SSIM_nuc\", \"pearson_mem\", \"SSIM_mem\"],\n",
    "    rot=30,\n",
    "    figsize=(8, 8),\n",
    ")\n",
    "plt.suptitle(\"Model Pixel Metrics\")\n",
    "plt.show()\n",
    "# Show boxplot of the metrics\n",
    "# Boxplot of the metrics\n",
    "test_segmentation_metrics.boxplot(\n",
    "    by=\"model\",\n",
    "    column=[\"jaccard\", \"accuracy\", \"mAP_75\",\"mAP_50\"],\n",
    "    rot=30,\n",
    "    figsize=(8, 8),\n",
    ")\n",
    "plt.suptitle(\"Model Segmentation Metrics\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ecdd41c",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "source": [
    "<div class=\"alert alert-info\">\n",
    "\n",
    "### Plotting the predictions and segmentations\n",
    "Here we will plot the predictions and segmentations side by side for the pretrained and trained models.<br>\n",
    "<ul>\n",
    "<li>How do yout model, the pretrained model and the ground truth compare?</li>\n",
    "<li>How do the segmentations compare? </li>\n",
    "</ul>\n",
    "Feel free to modify the crop size and Y,X slicing to view different areas of the FOV\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4fac45ba",
   "metadata": {
    "tags": []
   },
   "source": [
    "<div class=\"alert alert-success\">\n",
    "\n",
    "<h2> Checkpoint 2 </h2>\n",
    "\n",
    "Congratulations! You have completed the second checkpoint. You have:\n",
    "- Visualized the predictions and segmentations of the model. <br>\n",
    "- Evaluated the performance of the model using pixel-based metrics and segmentation-based metrics. <br>\n",
    "- Compared the performance of the model you trained with the pretrained model. <br>\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "482f1a41",
   "metadata": {
    "lines_to_next_cell": 0,
    "tags": []
   },
   "source": [
    "# Part 3: Visualizing the encoder and decoder features & exploring the model's range of validity\n",
    "\n",
    "- In this section, we will visualize the encoder and decoder features of the model you trained.\n",
    "- We will also explore the model's range of validity by looking at the feature maps of the encoder and decoder.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5eb7a5a",
   "metadata": {
    "tags": []
   },
   "source": [
    "<div class=\"alert alert-info\">\n",
    "<h3> Task 3.1: Let's look at what the model is learning </h3>\n",
    "\n",
    "- If you are unfamiliar with Principal Component Analysis (PCA), you can read up <a href=\"https://en.wikipedia.org/wiki/Principal_component_analysis\">here</a> <br>\n",
    "- Run the next cells. We will visualize the encoder feature maps of the trained model.\n",
    " We will use PCA to visualize the feature maps by mapping the first 3 principal components to a colormap `Color` <br>\n",
    "\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a663afa5",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Script to visualize the encoder feature maps of a trained model.\n",
    "Using PCA to visualize feature maps is inspired by\n",
    "https://doi.org/10.48550/arXiv.2304.07193 (Oquab et al., 2023).\n",
    "\"\"\"\n",
    "from typing import NamedTuple\n",
    "\n",
    "from matplotlib.patches import Rectangle\n",
    "from monai.networks.layers import GaussianFilter\n",
    "from skimage.exposure import rescale_intensity\n",
    "from skimage.transform import downscale_local_mean\n",
    "from sklearn.decomposition import PCA\n",
    "from sklearn.manifold import TSNE\n",
    "\n",
    "\n",
    "def feature_map_pca(feature_map: np.array, n_components: int = 8) -> PCA:\n",
    "    \"\"\"\n",
    "    Compute PCA on a feature map.\n",
    "    :param np.array feature_map: (C, H, W) feature map\n",
    "    :param int n_components: number of components to keep\n",
    "    :return: PCA: fit sklearn PCA object\n",
    "    \"\"\"\n",
    "    # (C, H, W) -> (C, H*W)\n",
    "    feat = feature_map.reshape(feature_map.shape[0], -1)\n",
    "    pca = PCA(n_components=n_components)\n",
    "    pca.fit(feat)\n",
    "    return pca\n",
    "\n",
    "def pcs_to_rgb(feat: np.ndarray, n_components: int = 8) -> np.ndarray:\n",
    "    pca = feature_map_pca(feat[0], n_components=n_components)\n",
    "    pc_first_3 = pca.components_[:3].reshape(3, *feat.shape[-2:])\n",
    "    return np.stack(\n",
    "        [rescale_intensity(pc, out_range=(0, 1)) for pc in pc_first_3], axis=-1\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87ec742b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the test dataset\n",
    "test_data_path = top_dir / \"06_image_translation/test/a549_hoechst_cellmask_test.zarr\"\n",
    "test_dataset = open_ome_zarr(test_data_path)\n",
    "\n",
    "# Looking at the test dataset\n",
    "print('Test dataset:')\n",
    "test_dataset.print_tree()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "abffa1a2",
   "metadata": {
    "lines_to_next_cell": 0,
    "tags": []
   },
   "source": [
    "<div class=\"alert alert-info\">\n",
    "\n",
    "- Change the `fov` and `crop` size to visualize the feature maps of the encoder and decoder  <br>\n",
    "Note: the crop should be a multiple of 384\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "056424d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load one position\n",
    "row = 0\n",
    "col = 0\n",
    "center_index = 2\n",
    "n = 1\n",
    "crop = 384 * n\n",
    "fov = 10\n",
    "\n",
    "# normalize phase\n",
    "norm_meta = test_dataset.zattrs[\"normalization\"][\"Phase3D\"][\"dataset_statistics\"]\n",
    "\n",
    "# Get the OME-Zarr metadata\n",
    "Y,X = test_dataset[f\"0/0/{fov}\"].data.shape[-2:]\n",
    "test_dataset.channel_names\n",
    "phase_idx= test_dataset.channel_names.index('Phase3D')\n",
    "assert crop//2 < Y and crop//2 < Y , \"Crop size larger than the image. Check the image shape\"\n",
    "\n",
    "phase_img = test_dataset[f\"0/0/{fov}/0\"][:, phase_idx:phase_idx+1,0:1, Y//2 - crop // 2 : Y//2 + crop // 2, X//2 - crop // 2 : X//2 + crop // 2]\n",
    "fluo = test_dataset[f\"0/0/{fov}/0\"][0, 1:3, 0, Y//2 - crop // 2 : Y//2 + crop // 2, X//2 - crop // 2 : X//2 + crop // 2]\n",
    "\n",
    "phase_img = (phase_img - norm_meta[\"median\"]) / norm_meta[\"iqr\"]\n",
    "plt.imshow(phase_img[0,0,0], cmap=\"gray\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "821c77b9",
   "metadata": {
    "lines_to_next_cell": 0,
    "tags": []
   },
   "source": [
    "<div class=\"alert alert-info\">\n",
    "For the following tasks we will use the pretrained model to extract the encoder and decoder features <br>\n",
    "Extra: If you are done with the whole checkpoint, you can try to look at what your trained model learned.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f52bb32",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Loading the pretrained model\n",
    "pretrained_model_ckpt = (\n",
    "    top_dir / \"06_image_translation/pretrained_models/VSCyto2D/epoch=399-step=23200.ckpt\"\n",
    ")\n",
    "# model config as before\n",
    "phase2fluor_config = dict(\n",
    "    in_channels=1,\n",
    "    out_channels=2,\n",
    "    encoder_blocks=[3, 3, 9, 3],\n",
    "    dims=[96, 192, 384, 768],\n",
    "    decoder_conv_blocks=2,\n",
    "    stem_kernel_size=(1, 2, 2),\n",
    "    in_stack_depth=1,\n",
    "    pretraining=False,\n",
    ")\n",
    "\n",
    "# load model\n",
    "model = VSUNet.load_from_checkpoint(\n",
    "    pretrained_model_ckpt,\n",
    "    architecture=\"UNeXt2_2D\",\n",
    "    model_config=phase2fluor_config.copy(),\n",
    "    accelerator=\"gpu\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3c1a56c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Extract features\n",
    "with torch.inference_mode():\n",
    "    # encoder\n",
    "    encoder_features = model.model.encoder(torch.from_numpy(phase_img.astype(np.float32)).to(model.device))[0]\n",
    "    encoder_features_np = [f.detach().cpu().numpy() for f in encoder_features]\n",
    "    \n",
    "    # Print the encoder features shapes\n",
    "    for f in encoder_features_np:\n",
    "        print(f.shape)\n",
    "\n",
    "    # decoder\n",
    "    features = encoder_features.copy()\n",
    "    features.reverse()\n",
    "    feat = features[0]\n",
    "    features.append(None)\n",
    "    decoder_features_np = []\n",
    "    for skip, stage in zip(features[1:], model.model.decoder.decoder_stages):\n",
    "        feat = stage(feat, skip)\n",
    "        decoder_features_np.append(feat.detach().cpu().numpy())\n",
    "    for f in decoder_features_np:\n",
    "        print(f.shape)\n",
    "    prediction = model.model.head(feat).detach().cpu().numpy()\n",
    "    \n",
    "# Defining the colors for plotting\n",
    "class Color(NamedTuple):\n",
    "    r: float\n",
    "    g: float\n",
    "    b: float\n",
    "# Defining the colors for plottting the PCA\n",
    "BOP_ORANGE = Color(0.972549, 0.6784314, 0.1254902)\n",
    "BOP_BLUE = Color(BOP_ORANGE.b, BOP_ORANGE.g, BOP_ORANGE.r)\n",
    "GREEN = Color(0.0, 1.0, 0.0)\n",
    "MAGENTA = Color(1.0, 0.0, 1.0)\n",
    "\n",
    "# Defining the functions to rescale the image and composite the nuclear and membrane images\n",
    "def rescale_clip(image: torch.Tensor) -> np.ndarray:\n",
    "    return rescale_intensity(image, out_range=(0, 1))[..., None].repeat(3, axis=-1)\n",
    "\n",
    "def composite_nuc_mem(\n",
    "    image: torch.Tensor, nuc_color: Color, mem_color: Color\n",
    ") -> np.ndarray:\n",
    "    c_nuc = rescale_clip(image[0]) * nuc_color\n",
    "    c_mem = rescale_clip(image[1]) * mem_color\n",
    "    return rescale_intensity(c_nuc + c_mem, out_range=(0, 1))\n",
    "\n",
    "def clip_p(image: np.ndarray) -> np.ndarray:\n",
    "    return rescale_intensity(image.clip(*np.percentile(image, [1, 99])))\n",
    "\n",
    "def clip_highlight(image: np.ndarray) -> np.ndarray:\n",
    "    return rescale_intensity(image.clip(0, np.percentile(image, 99.5)))\n",
    "\n",
    "# Plot the PCA to RGB of the feature maps\n",
    "f, ax = plt.subplots(10, 1, figsize=(5, 25))\n",
    "n_components = 4\n",
    "ax[0].imshow(phase_img[0, 0, 0], cmap=\"gray\")\n",
    "ax[0].set_title(f\"Phase {phase_img.shape[1:]}\")\n",
    "ax[-1].imshow(clip_p(composite_nuc_mem(fluo, GREEN, MAGENTA)))\n",
    "ax[-1].set_title(\"Fluorescence\")\n",
    "\n",
    "for level, feat in enumerate(encoder_features_np):\n",
    "    ax[level + 1].imshow(pcs_to_rgb(feat, n_components=n_components))\n",
    "    ax[level + 1].set_title(f\"Encoder stage {level+1} {feat.shape[1:]}\")\n",
    "\n",
    "for level, feat in enumerate(decoder_features_np):\n",
    "    ax[5 + level].imshow(pcs_to_rgb(feat, n_components=n_components))\n",
    "    ax[5 + level].set_title(f\"Decoder stage {level+1} {feat.shape[1:]}\")\n",
    "\n",
    "pred_comp = composite_nuc_mem(prediction[0, :, 0], BOP_BLUE, BOP_ORANGE)\n",
    "ax[-2].imshow(clip_p(pred_comp))\n",
    "ax[-2].set_title(f\"Prediction {prediction.shape[1:]}\")\n",
    "\n",
    "for a in ax.ravel():\n",
    "    a.axis(\"off\")\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bbf924ed",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [],
   "source": [
    "YX_PATCH_SIZE = (256*2,256*2)\n",
    "source_channel = [\"Phase3D\"]\n",
    "target_channel = [\"Nucl\", \"Mem\"]\n",
    "\n",
    "normalizations = [\n",
    "    NormalizeSampled(\n",
    "        keys=source_channel,\n",
    "        level=\"fov_statistics\",\n",
    "        subtrahend=\"mean\",\n",
    "        divisor=\"std\",\n",
    "    ),\n",
    "        NormalizeSampled(\n",
    "        keys=target_channel,\n",
    "        level=\"fov_statistics\",\n",
    "        subtrahend=\"median\",\n",
    "        divisor=\"iqr\",\n",
    "    )\n",
    "]\n",
    "\n",
    "# Re-load the dataloader\n",
    "phase2fluor_2D_data = HCSDataModule(\n",
    "    data_path,\n",
    "    architecture=\"UNeXt2_2D\",\n",
    "    source_channel=source_channel,\n",
    "    target_channel=target_channel,\n",
    "    z_window_size=1,\n",
    "    split_ratio=0.8,\n",
    "    batch_size=1,\n",
    "    num_workers=8,\n",
    "    yx_patch_size=YX_PATCH_SIZE,\n",
    "    augmentations=[],\n",
    "    normalizations=normalizations,\n",
    ")\n",
    "phase2fluor_2D_data.setup(\"test\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ba7ed50",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# ########## TODO ##############\n",
    "batch_number = 3 # Change this to see different batches of data\n",
    "# #######################\n",
    "y_slice = slice(Y//2-256*n//2, Y//2+256*n//2)\n",
    "x_slice = slice(X//2-256*n//2, X//2+256*n//2)\n",
    "\n",
    "# Iterate through the test dataloader to get the desired batch\n",
    "i = 0\n",
    "for batch in phase2fluor_2D_data.test_dataloader():\n",
    "    # break if we reach the desired batch    \n",
    "    if i == batch_number-1:\n",
    "        break\n",
    "    i += 1\n",
    "\n",
    "# Plot the batch source and target images\n",
    "f, ax = plt.subplots(1, 2, figsize=(8, 12))\n",
    "target_composite = composite_nuc_mem(batch[\"target\"][0].cpu().numpy(), GREEN, MAGENTA)\n",
    "ax[0].imshow(\n",
    "    batch[\"source\"][0, 0, 0,y_slice,x_slice].cpu().numpy(), cmap=\"gray\", vmin=-15, vmax=15\n",
    ")\n",
    "ax[1].imshow(clip_highlight(target_composite[0,y_slice,x_slice]))\n",
    "for a in ax.ravel():\n",
    "    a.axis(\"off\")\n",
    "f.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42ae3119",
   "metadata": {
    "lines_to_next_cell": 0,
    "tags": []
   },
   "source": [
    "<div class=\"alert alert-info\">\n",
    "\n",
    "### Task 3.3: Using the selected batch to test the model's range of validity\n",
    "\n",
    "- Given the selected batch use `monai.networks.layers.GaussianFilter` to blur the images with different sigmas.\n",
    " Check the documentation <a href=\"https://docs.monai.io/en/stable/networks.html#gaussianfilter\">here</a> <br>\n",
    "- Plot the source and predicted images comparing the source, target and added perturbations <br>\n",
    "- How is the model's predictions given the perturbations? <br>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de4ebac7",
   "metadata": {
    "tags": [
     "solution"
    ]
   },
   "outputs": [],
   "source": [
    "# ########## SOLUTION ##############\n",
    "# Try out different multiples of 256 to visualize larger/smaller crops\n",
    "n = 3\n",
    "# ##############################\n",
    "# Center cropping the image\n",
    "y_slice = slice(Y//2-256*n//2, Y//2+256*n//2)\n",
    "x_slice = slice(X//2-256*n//2, X//2+256*n//2)\n",
    "\n",
    "f, ax = plt.subplots(3, 2, figsize=(8, 12))\n",
    "\n",
    "target_composite = composite_nuc_mem(batch[\"target\"][0].cpu().numpy(), GREEN, MAGENTA)\n",
    "ax[0, 0].imshow(\n",
    "    batch[\"source\"][0, 0, 0,y_slice,x_slice].cpu().numpy(), cmap=\"gray\", vmin=-15, vmax=15\n",
    ")\n",
    "ax[0, 1].imshow(clip_highlight(target_composite[0,y_slice,x_slice]))\n",
    "ax[0,0].set_title('Source and target')\n",
    "\n",
    "# no perturbation\n",
    "with torch.inference_mode():\n",
    "    phase = batch[\"source\"].to(model.device)[:,:,:,y_slice,x_slice]\n",
    "    pred = model(phase).cpu().numpy()\n",
    "pred_composite = composite_nuc_mem(pred[0], BOP_BLUE, BOP_ORANGE)\n",
    "ax[1, 0].imshow(phase[0,0,0].cpu().numpy(), cmap=\"gray\", vmin=-15, vmax=15)\n",
    "ax[1, 1].imshow(pred_composite[0])\n",
    "ax[1,0].set_title('No perturbation')\n",
    "\n",
    "\n",
    "# Select a sigma for the Gaussian filtering\n",
    "# ########## SOLUTION ##############\n",
    "# Tensor dimensions (B,C,D,H,W). \n",
    "# Hint: Use the GaussianFilter layer to blur the phase image. Provide the  num spatial dimensions and sigma\n",
    "# Hint: Spatial (D,H,W). Apply the same sigma to H,W\n",
    "gaussian_blur = GaussianFilter(spatial_dims=3, sigma=(0,2,2))\n",
    "# #############################\n",
    "with torch.inference_mode():\n",
    "    phase = batch[\"source\"].to(model.device)[:,:,:,y_slice,x_slice]\n",
    "    phase = gaussian_blur(phase)\n",
    "    pred = model(phase).cpu().numpy()\n",
    "pred_composite = composite_nuc_mem(pred[0], BOP_BLUE, BOP_ORANGE)\n",
    "ax[2, 0].imshow(phase[0, 0, 0].cpu().numpy(), cmap=\"gray\", vmin=-15, vmax=15)\n",
    "ax[2, 1].imshow(pred_composite[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe8b5a8b",
   "metadata": {
    "tags": []
   },
   "source": [
    "<div class=\"alert alert-info\">\n",
    "\n",
    "### Task 3.3: Using the selected batch to test the model's range of validity\n",
    "\n",
    "- Scale the pixel values up/down of the phase image <br>\n",
    "- Plot the source and predicted images comparing the source, target and added perturbations <br>\n",
    "- How is the model's predictions given the perturbations? <br>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24978c2b",
   "metadata": {
    "tags": [
     "solution"
    ]
   },
   "outputs": [],
   "source": [
    "n = 3\n",
    "y_slice = slice(Y//2, Y//2+256*n)\n",
    "x_slice = slice(X//2, X//2+256*n)\n",
    "f, ax = plt.subplots(3, 2, figsize=(8, 12))\n",
    "\n",
    "target_composite = composite_nuc_mem(batch[\"target\"][0].cpu().numpy(), GREEN, MAGENTA)\n",
    "ax[0, 0].imshow(\n",
    "    batch[\"source\"][0, 0, 0,y_slice,x_slice].cpu().numpy(), cmap=\"gray\", vmin=-15, vmax=15\n",
    ")\n",
    "ax[0, 1].imshow(clip_highlight(target_composite[0,y_slice,x_slice]))\n",
    "ax[0,0].set_title('Source and target')\n",
    "\n",
    "# no perturbation\n",
    "with torch.inference_mode():\n",
    "    phase = batch[\"source\"].to(model.device)[:,:,:,y_slice,x_slice]\n",
    "    pred = model(phase).cpu().numpy()\n",
    "pred_composite = composite_nuc_mem(pred[0], BOP_BLUE, BOP_ORANGE)\n",
    "ax[1, 0].imshow(phase[0,0,0].cpu().numpy(), cmap=\"gray\", vmin=-15, vmax=15)\n",
    "ax[1, 1].imshow(pred_composite[0])\n",
    "ax[1,0].set_title('No perturbation')\n",
    "\n",
    "\n",
    "# Rescale the pixel value up/down\n",
    "with torch.inference_mode():\n",
    "    phase = batch[\"source\"].to(model.device)[:,:,:,y_slice,x_slice]\n",
    "    # ########## SOLUTION ##############\n",
    "    # Hint: Scale the phase intensity up/down until the model breaks\n",
    "    phase = phase * 10\n",
    "    # #######################\n",
    "    pred = model(phase).cpu().numpy()\n",
    "pred_composite = composite_nuc_mem(pred[0], BOP_BLUE, BOP_ORANGE)\n",
    "ax[2, 0].imshow(phase[0, 0, 0].cpu().numpy(), cmap=\"gray\", vmin=-15, vmax=15)\n",
    "ax[2, 1].imshow(pred_composite[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e8b35502",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-warning\">\n",
    "<h3> Questions </h3>\n",
    "How is the model's predictions given the blurring and scaling perturbations? <br>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8868857",
   "metadata": {
    "lines_to_next_cell": 0,
    "tags": [
     "solution"
    ]
   },
   "outputs": [],
   "source": [
    "# ########## SOLUTIONS FOR ALL POSSIBLE PLOTTINGS ##############\n",
    "# This plots all perturbations\n",
    "\n",
    "n = 3\n",
    "y_slice = slice(Y//2, Y//2+256*n)\n",
    "x_slice = slice(X//2, X//2+256*n)\n",
    "f, ax = plt.subplots(6, 2, figsize=(8, 12))\n",
    "\n",
    "target_composite = composite_nuc_mem(batch[\"target\"][0].cpu().numpy(), GREEN, MAGENTA)\n",
    "ax[0, 0].imshow(\n",
    "    batch[\"source\"][0, 0, 0,y_slice,x_slice].cpu().numpy(), cmap=\"gray\", vmin=-15, vmax=15\n",
    ")\n",
    "ax[0, 1].imshow(clip_highlight(target_composite[0,y_slice,x_slice]))\n",
    "ax[0,0].set_title('Source and target')\n",
    "\n",
    "# no perturbation\n",
    "with torch.inference_mode():\n",
    "    phase = batch[\"source\"].to(model.device)[:,:,:,y_slice,x_slice]\n",
    "    pred = model(phase).cpu().numpy()\n",
    "pred_composite = composite_nuc_mem(pred[0], BOP_BLUE, BOP_ORANGE)\n",
    "ax[1, 0].imshow(phase[0,0,0].cpu().numpy(), cmap=\"gray\", vmin=-15, vmax=15)\n",
    "ax[1, 1].imshow(pred_composite[0])\n",
    "ax[1,0].set_title('No perturbation')\n",
    "\n",
    "\n",
    "# 2-sigma gaussian blur\n",
    "gaussian_blur = GaussianFilter(spatial_dims=3, sigma=(0,2,2))\n",
    "with torch.inference_mode():\n",
    "    phase = batch[\"source\"].to(model.device)[:,:,:,y_slice,x_slice]\n",
    "    phase = gaussian_blur(phase)\n",
    "    pred = model(phase).cpu().numpy()\n",
    "pred_composite = composite_nuc_mem(pred[0], BOP_BLUE, BOP_ORANGE)\n",
    "ax[2, 0].imshow(phase[0, 0, 0].cpu().numpy(), cmap=\"gray\", vmin=-15, vmax=15)\n",
    "ax[2, 1].imshow(pred_composite[0])\n",
    "ax[2,0].set_title('Gaussian Blur Sigma=2')\n",
    "\n",
    "\n",
    "# 5-sigma gaussian blur\n",
    "gaussian_blur = GaussianFilter(spatial_dims=3, sigma=(0,5,5))\n",
    "with torch.inference_mode():\n",
    "    phase = batch[\"source\"].to(model.device)[:,:,:,y_slice,x_slice]\n",
    "    phase = gaussian_blur(phase)\n",
    "    pred = model(phase).cpu().numpy()\n",
    "pred_composite = composite_nuc_mem(pred[0], BOP_BLUE, BOP_ORANGE)\n",
    "ax[3, 0].imshow(phase[0, 0, 0].cpu().numpy(), cmap=\"gray\", vmin=-15, vmax=15)\n",
    "ax[3, 1].imshow(pred_composite[0])\n",
    "ax[3,0].set_title('Gaussian Blur Sigma=5')\n",
    "\n",
    "\n",
    "# 0.1x scaling\n",
    "with torch.inference_mode():\n",
    "    phase = batch[\"source\"].to(model.device)[:,:,:,y_slice,x_slice]\n",
    "    phase = phase*0.1\n",
    "    pred = model(phase).cpu().numpy()\n",
    "pred_composite = composite_nuc_mem(pred[0], BOP_BLUE, BOP_ORANGE)\n",
    "ax[4, 0].imshow(phase[0, 0, 0].cpu().numpy(), cmap=\"gray\", vmin=-15, vmax=15)\n",
    "ax[4, 1].imshow(pred_composite[0])\n",
    "ax[4,0].set_title('0.1x scaling')\n",
    "\n",
    "# 10x scaling\n",
    "with torch.inference_mode():\n",
    "    phase = batch[\"source\"].to(model.device)[:,:,:,y_slice,x_slice]\n",
    "    phase = phase* 10 \n",
    "    pred = model(phase).cpu().numpy()\n",
    "pred_composite = composite_nuc_mem(pred[0], BOP_BLUE, BOP_ORANGE)\n",
    "ax[5, 0].imshow(phase[0, 0, 0].cpu().numpy(), cmap=\"gray\", vmin=-15, vmax=15)\n",
    "ax[5, 1].imshow(pred_composite[0])\n",
    "ax[5,0].set_title('10x scaling')\n",
    "\n",
    "for a in ax.ravel():\n",
    "    a.axis(\"off\")\n",
    "\n",
    "f.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c822491",
   "metadata": {
    "tags": []
   },
   "source": [
    "<div class=\"alert alert-success\">\n",
    "\n",
    "<h2>\n",
    "🎉 The end of the notebook 🎉\n",
    "</h2>\n",
    "\n",
    "Congratulations! You have trained an image translation model, evaluated its performance, and explored what the network has learned. \n",
    "\n",
    "</div>"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "cell_metadata_filter": "all",
   "main_language": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}