exercise.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "95227378-32db-418d-aa85-bf29eb8ad145",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Instance Segmentation\n",
    "\n",
    "<div class=\"alert alert-block alert-danger\">Set your kernel to `06_instance_segmentation`\n",
    "</div>\n",
    "\n",
    "- So far we were only interested in semantic classes, eg foreground / background, cell types, person / car, etc. But in many cases we not only want to know if a certain pixel belongs to an object, but also to **which** unique object.\n",
    "\n",
    "\n",
    "- For isolated objects, this is trivial, all connected foreground pixels form one instance, yet often instances are very close together or even overlapping. Then we need to think a bit more how to formulate the inputs / loss to our network and how to extract the instances from the predictions.\n",
    "\n",
    "\n",
    "- Below is an example of the differences between object detection, semantic segmentation, and instance segmentation. The raw data shows a 2d slice of neural tissue acquired from a high resolution electron microscope. We could teach a network to simply detect each mitochondria (object detection). If we want to assign every pixel to a specific class, we could do semantic segmentation (like in the previous exercise). In this case we would have two classes, one class for neurons that contain mitochondria (green) and another class for neurons which do not. Finally, we could assign every pixel as belonging to a unique object (neurons in this case). This is instance segmentation and is an approach that can be very useful for biological data.\n",
    "\n",
    "![example_image](static/instance_seg.png)\n",
    "\n",
    "\n",
    "- If you are running this in jupyter lab, every markdown header is collapsible and all cells are collapsed by default. Just click on the left of a cell to expand it, and just make sure to expand until the code cells show. The headings unfortunately do not collapse in jupyter notebook, but will still give you an idea of breaks between exercises. \n",
    "\n",
    "\n",
    "- Most TODOs build off the previous TODOs and require copying over classes/functions before adding more to them. It could be annoying to constantly scroll back and forth so you can click the table of contents on the left hand side to easily navigate between sections.\n",
    "\n",
    "- If you have questions please let us know. Have fun!!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "230f4cec",
   "metadata": {
    "tags": []
   },
   "source": [
    "<div class=\"alert alert-block alert-success\"><h2>Start here (AKA checkpoint 0)</h1>\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8812ec06-d766-4809-8d5f-aade999bbbc4",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "## Importing packages\n",
    "\n",
    "* You should have already set up your conda environment by now. Import these packages and let us know if something fails so we can debug before moving on. \n",
    "* We're also importing a UNet model from the corresponding module, and other predefined functions from the `utils.py`. Please check them out once you'll be using them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17351d30-450d-4a86-96f3-999cdc5c1225",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import os\n",
    "import random\n",
    "import torch\n",
    "import datetime\n",
    "\n",
    "from glob import glob\n",
    "from natsort import natsorted\n",
    "from torch.utils.data import DataLoader\n",
    "from torch.utils.data import Dataset\n",
    "\n",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "from torchsummary import summary\n",
    "from tqdm import tqdm\n",
    "\n",
    "torch.backends.cudnn.benchmark = True"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c7bedd4-985d-41a9-b4c6-7c5669b716f0",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Part 1: Creating a simple model\n",
    "\n",
    "- Let's start by creating a simple model similar to the one in the semantic segmentation exercise. We will then improve it in the subsequent checkpoints"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "172f39f0-5270-4d48-a42e-61029d99ea72",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Load and visualize data\n",
    "    \n",
    "    \n",
    "- For this exercise we will be using data from [TissueNet](https://datasets.deepcell.org/) (paper [here](https://www.ncbi.nlm.nih.gov/pmc/articles/PMC9010346/)).\n",
    "    \n",
    "    \n",
    "- For our purposes we will use 50 training images and 20 testing images. The data is stored as tifs using the following structure:\n",
    "```\n",
    "woodshole/\n",
    "    ├── test\n",
    "    │   ├── img_0_cyto_masks.tif\n",
    "    │   ├── img_0_nuclei_masks.tif\n",
    "    │   ├── img_0.tif\n",
    "    │   ├── img_1_cyto_masks.tif\n",
    "    │   ├── img_1_nuclei_masks.tif\n",
    "    │   └── img_1.tif\n",
    "    │  \n",
    "    └── train\n",
    "        ├── img_0_cyto_masks.tif\n",
    "        ├── img_0_nuclei_masks.tif\n",
    "        ├── img_0.tif\n",
    "        ├── img_1_cyto_masks.tif\n",
    "        ├── img_1_nuclei_masks.tif\n",
    "        └── img_1.tif\n",
    "```\n",
    "- Each raw image is stored as `img_{n}.tif` and is already stored as float32 between 0 and 1 so does not need to be normalized for training purposes. There are two channels in the raw data, one for nuclei and one for cytoplasm. \n",
    "    \n",
    "    \n",
    "- The corresponding mask files contain instance segmentations for the raw data. The nuclei masks correspond to the first channel of the raw data. The cytoplasm masks correspond to nucleus + cytoplasm. We will start with the nuclei data, as it is likely easier to segment than the cytoplasm, but you will be able to apply the techniques you learn on the harder data at the end of the exercise, time permitting.\n",
    "\n",
    "![example_image](static/example_image.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "820db827-198b-465b-941c-22d3452f3664",
   "metadata": {},
   "outputs": [],
   "source": [
    "# lets start by loading our images into lists so we can visualize and get oriented with our data\n",
    "# natsorted is a package that takes away some of the annoyances of the regular sorting function\n",
    "# glob is a package that allows us to load files from directories. Feel free to inspect these lists.\n",
    "\n",
    "train_cyto = natsorted(glob('woodshole/train/*cyto*'))\n",
    "train_nuclei = natsorted(glob('woodshole/train/*nuclei*'))\n",
    "train_raw = [i for i in natsorted(glob('woodshole/train/*.tif')) if 's.tif' not in i]\n",
    "\n",
    "test_cyto = natsorted(glob('woodshole/test/*cyto*'))\n",
    "test_nuclei = natsorted(glob('woodshole/test/*nuclei*'))\n",
    "test_raw = [i for i in natsorted(glob('woodshole/test/*.tif')) if 's.tif' not in i]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "223540b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's define the device we'll be using throughout the notebook\n",
    "\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9cfdebc1-b8d1-4ed9-9e71-cb3ec35305cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# convenience functions for viewing labels as rgb, and reading files into numpy arrays\n",
    "from skimage import color\n",
    "from skimage.io import imread\n",
    "\n",
    "# utility function to view labels as rgb lut with matplotlib\n",
    "# eg plt.imshow(create_lut(labels))\n",
    "from utils import create_lut"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb7fbe07-8be2-426f-b256-cc8a20e48b20",
   "metadata": {},
   "outputs": [],
   "source": [
    "# select a random cytoplasm mask file\n",
    "cyto_file = random.choice(train_cyto)\n",
    "\n",
    "# use skimage.io.imread to read our data into numpy arrays\n",
    "cyto = imread(cyto_file)\n",
    "nuclei = imread(cyto_file.replace('cyto', 'nuclei'))\n",
    "raw = imread(cyto_file.replace('_cyto_masks', ''))\n",
    "\n",
    "#our raw data shape is (c, h, w) and there are only two channels.\n",
    "#to visualize as an rgb image we need to add another dummy dimension and then transpose so that it is (h,w,3)\n",
    "raw = np.vstack((raw, np.zeros_like(raw)[:1]))\n",
    "raw = raw.transpose(1,2,0)\n",
    "\n",
    "# visualize the data - execute this cell a few times to see different examples.\n",
    "# you can also change train_cyto to test_cyto above to see some test data. it is pretty similar \n",
    "fig, axes = plt.subplots(1,5,figsize=(20, 10),sharex=True,sharey=True,squeeze=False)\n",
    "\n",
    "axes[0][0].imshow(raw[:,:,0], cmap='gray')\n",
    "axes[0][0].title.set_text('Raw nuclei channel')\n",
    "\n",
    "axes[0][1].imshow(raw[:,:,1], cmap='gray')\n",
    "axes[0][1].title.set_text('Raw cyto channel')\n",
    "\n",
    "axes[0][2].imshow(raw)\n",
    "axes[0][2].title.set_text('Raw overlay')\n",
    "\n",
    "axes[0][3].imshow(raw[:,:,0], cmap='gray')\n",
    "axes[0][3].imshow(create_lut(nuclei), alpha=0.5)\n",
    "axes[0][3].title.set_text('nuclei mask')\n",
    "\n",
    "axes[0][4].imshow(raw[:,:,1], cmap='gray')\n",
    "axes[0][4].imshow(create_lut(cyto), alpha=0.5)\n",
    "axes[0][4].title.set_text('cyto mask')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00ea412d-b780-4028-a963-3012a29b6127",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Create simple augmentation function</h3>\n",
    "    \n",
    "- You have already learned about the importance of augmenting your training data.\n",
    "- For our exercise we will use an augmentation library called [albumentations](https://albumentations.ai/) which provides easy to use, fast transforms\n",
    "- Here is a nice tutorial: https://albumentations.ai/docs/examples/example_kaggle_salt/\n",
    "- To start, we will add a few simple augmentations to both our raw and mask data:\n",
    "    - randomly crop a 64x64 patch\n",
    "    - horizontally flip with a 50% probability\n",
    "    - vertically flip with a 50% probability\n",
    "- Below is a simple example of a random crop augmentation. \n",
    "\n",
    "![example_augmentation](static/example_augmentation.png)\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7815224",
   "metadata": {},
   "outputs": [],
   "source": [
    "# import albumentations library\n",
    "import albumentations as A\n",
    "\n",
    "file = random.choice(train_nuclei)\n",
    "\n",
    "full_mask_nuclei = imread(file)\n",
    "full_raw_nuclei = imread(file.replace('_nuclei_masks', ''))[0]\n",
    "\n",
    "# Define simple augmentation pipeline\n",
    "transform = A.Compose([\n",
    "              A.RandomCrop(width=64, height=64),\n",
    "              A.HorizontalFlip(p=0.5),\n",
    "              A.VerticalFlip(p=0.5)\n",
    "            ])\n",
    "\n",
    "transformed = transform(image=full_raw_nuclei, mask=full_mask_nuclei)\n",
    "          \n",
    "aug_raw, aug_mask = transformed['image'], transformed['mask']\n",
    "\n",
    "fig, axes = plt.subplots(1,2,figsize=(10, 10),sharex=False,sharey=False,squeeze=False)\n",
    "\n",
    "axes[0][0].imshow(full_raw_nuclei, cmap='gray')\n",
    "axes[0][0].imshow(create_lut(full_mask_nuclei), alpha=0.5)\n",
    "axes[0][0].set_title('original image')\n",
    "\n",
    "axes[0][1].imshow(aug_raw, cmap='gray')\n",
    "axes[0][1].imshow(create_lut(aug_mask), alpha=0.5)\n",
    "axes[0][1].set_title('example random augmentation');"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2aec8caf",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Create fg/bg representation\n",
    "    \n",
    "- It would be ideal to directly predict unique labels in a dataset. Unfortunately this requires global information which can become difficult as datasets increase in size. Consequently, alternative approaches aim to solve the problem locally.\n",
    "    \n",
    "    \n",
    "- We will start with the most trivial approach: learning a foreground / background mask and then relabeling connected pixels as unique objects. While this approach might suffice on simple datasets, you will see how it can become problematic on datasets in which objects are tightly packed.\n",
    "\n",
    "![example_fgbg](static/example_fgbg.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab58aa51-348c-4802-87aa-7f40a57c9eb6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# import utility function to erode borders\n",
    "from utils import erode\n",
    "\n",
    "# visualize the representation (repeatedly run cell)\n",
    "\n",
    "file = random.choice(train_nuclei)\n",
    "\n",
    "full_mask_nuclei = imread(file)\n",
    "full_raw_nuclei = imread(file.replace('_nuclei_masks', ''))[0]\n",
    "\n",
    "# Define simple augmentation pipeline\n",
    "transform = A.Compose([\n",
    "              A.RandomCrop(width=64, height=64),\n",
    "              A.HorizontalFlip(p=0.5),\n",
    "              A.VerticalFlip(p=0.5)\n",
    "            ])\n",
    "\n",
    "transformed = transform(image=full_raw_nuclei, mask=full_mask_nuclei)\n",
    "          \n",
    "aug_raw, aug_mask = transformed['image'], transformed['mask']\n",
    "\n",
    "# erode label borders\n",
    "eroded_labels = erode(\n",
    "    aug_mask,\n",
    "    iterations=1,\n",
    "    border_value=1)\n",
    "\n",
    "# create fg/bg mask\n",
    "labels_two_class = (eroded_labels != 0).astype(np.float32)\n",
    "\n",
    "# check num classes and pixel counts\n",
    "print(np.unique(labels_two_class, return_counts=True))\n",
    "\n",
    "fig, axes = plt.subplots(1,3,figsize=(15, 15),sharex=True,sharey=True,squeeze=False)\n",
    "\n",
    "axes[0][0].imshow(aug_raw, cmap='gray')\n",
    "axes[0][0].title.set_text('Raw')\n",
    "\n",
    "axes[0][0].imshow(create_lut(aug_mask), alpha=0.5)\n",
    "axes[0][0].title.set_text('Labels')\n",
    "\n",
    "axes[0][1].imshow(aug_raw, cmap='gray')\n",
    "axes[0][1].title.set_text('Raw')\n",
    "\n",
    "axes[0][1].imshow(create_lut(eroded_labels), alpha=0.5)\n",
    "axes[0][1].title.set_text('Eroded labels')\n",
    "\n",
    "axes[0][2].imshow(labels_two_class)\n",
    "axes[0][2].title.set_text('Foreground / background')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02c3ec11",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Create simple dataset\n",
    "\n",
    "- Now let's combine all this into a simple dataset (similar to what was done in the image segmentation exercise)\n",
    "- For now our dataset should just load the raw and mask data, and apply simple augmentation.\n",
    "- We will just use the first channel of the raw data (nuclei) for now (`raw[0]`)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f61492b1",
   "metadata": {
    "tags": []
   },
   "source": [
    "<a id='first-todo'></a>\n",
    "<div class=\"alert alert-block alert-info\"><h4> Task 1.1</h4>\n",
    "\n",
    "* Create a basic dataset with simple data augmentations called `TissueNetDataset`\n",
    "* The `__init__` method should load our **mask file names** and **raw file names** into sorted lists\n",
    "* Add a parameter to your `TissueNetDataset` to define a `crop_size` of `64 pixels`, that will be used in the `augment_data` method\n",
    "* Add a boolean parameter to your `TissueNetDataset` to create a validation split `val_split`\n",
    "* The `augment_data` method should take in a raw and mask array and return the augmented raw and mask arrays.\n",
    "* In the `__getitem__` method, we should only augment our data if our split is `train`. So you will need to also add split as an attribute in the `__init__` method\n",
    "* After augmenting your data, create a fg / bg representation\n",
    "* Erode your fg / bg representation by a pixel to introduce a seperation or explicit boundary between touching labels\n",
    "* Make sure to return your fg/bg as float32 for training\n",
    "* Make sure to add a dummy channel dimension to your arrays for training (Pytorch assumes we have tensor shape batch, channel, height, width). We will add a batch dimension later once we create a data loader.\n",
    "* For now we will just use the nuclei channel for training, so make sure to slice the correct channel of the raw data before returning\n",
    "* You should return the raw and mask arrays\n",
    "* If you want to learn more about the ETL(extract, transform, load) pipeline in torch, click [here](https://pytorch.org/docs/stable/data.html)\n",
    "* Click [here](#second-todo) if you need to go to the next **Task**\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c3829fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TissueNetDataset(Dataset):\n",
    "    def __init__(self,\n",
    "                 root_dir,\n",
    "                 split='train',\n",
    "                 mask='nuclei',\n",
    "                 crop_size=None,\n",
    "                 val_split=False\n",
    "                ):\n",
    "        \n",
    "        # make sure to add your split method since we will use it in `__getitem__`\n",
    "        self.split = ...\n",
    "        # make sure to add your crop size\n",
    "        self.crop_size = ...\n",
    "\n",
    "        # using root_dir, split and mask create a path to files and sort it \n",
    "        # Hint: natsorted glob and os libraries could come in handy\n",
    "        self.mask_files = ... # load mask files into sorted list\n",
    "        self.raw_files = ... # load image files into sorted list\n",
    "        \n",
    "        # Add another parameter `val_split`. \n",
    "        # **There are 20 test files in total**\n",
    "        # If `split` is `test` and `val_split` is True, take the first half of the test files\n",
    "        # else take the second half\n",
    "        \n",
    "        \n",
    "    def __len__(self):\n",
    "        return len(self.raw_files)\n",
    "    \n",
    "    def augment_data(self, raw, mask):\n",
    "        \n",
    "        transform = ... # create your augmentations\n",
    "        \n",
    "        transformed = ... # call your augmentations\n",
    "        \n",
    "        raw, mask = ... # get your resulting arrays\n",
    "        \n",
    "        return raw, mask\n",
    "        \n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        raw_file = ... # get raw file at index\n",
    "        mask_file = ... # get mask file at index\n",
    "        \n",
    "        raw = ... # load raw to numpy array\n",
    "        mask = ... # load mask to numpy array\n",
    "        \n",
    "        raw = ... # get nuclei channel\n",
    "        \n",
    "        # augment your data if split mode is train\n",
    "        \n",
    "        mask = ... # erode your labels, cast to float32.\n",
    "\n",
    "        raw = ... # add channel dimension to comply with pytorch standard (C, H, W)\n",
    "        mask = ... # add channel dimension\n",
    "            \n",
    "        return raw, mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76c1cc8d",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = TissueNetDataset(root_dir='woodshole', split='train', crop_size=64)\n",
    "test_dataset = TissueNetDataset(root_dir='woodshole', split='test')\n",
    "val_dataset = TissueNetDataset(root_dir='woodshole', split='test', val_split=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9450189a-647f-43e4-bda1-62e203653186",
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils import relabel_cc\n",
    "# visualize the representation (repeatedly run cell)\n",
    "\n",
    "raw, mask = train_dataset[random.randrange(len(train_dataset))]\n",
    "\n",
    "labels_two_class = (mask != 0).astype(np.float32)\n",
    "\n",
    "fig, axes = plt.subplots(1,2,figsize=(10, 10),sharex=True,sharey=True,squeeze=False)\n",
    "\n",
    "axes[0][0].imshow(raw.squeeze(), cmap='gray')\n",
    "axes[0][0].title.set_text('Raw')\n",
    "\n",
    "axes[0][0].imshow(create_lut(relabel_cc(mask.squeeze().astype(int))), alpha=0.5)\n",
    "axes[0][0].title.set_text('Segmentation')\n",
    "\n",
    "axes[0][1].imshow(labels_two_class.squeeze())\n",
    "axes[0][1].title.set_text('Foreground / background')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4baf376-337e-4906-abde-7c65d7a3c8d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# run cell repeatedly to see different crops\n",
    "\n",
    "raw, mask = train_dataset[random.randrange(len(train_dataset))]\n",
    "\n",
    "fig, axes = plt.subplots(1,2,figsize=(10, 10),sharex=True,sharey=True,squeeze=False)\n",
    "\n",
    "axes[0][0].imshow(raw.squeeze(), cmap='gray')\n",
    "axes[0][0].title.set_text('Raw')\n",
    "\n",
    "axes[0][1].imshow(mask.squeeze())\n",
    "axes[0][1].title.set_text('Mask')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "594c20e0-4cc6-419f-a1d8-8290ba39aa9e",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Create shallow network, visualize receptive field\n",
    "    \n",
    "- Let's create a shallow two level U-Net and visualize the receptive field. We will see later how this receptive field changes as we add more levels and change our input image size\n",
    "    \n",
    "    \n",
    "- The receptive field tells us how much of the image the network is looking at in each level -- this is the amount of spatial context that the network can use to create predictions.\n",
    "    \n",
    "    \n",
    "- Run the following cell to see the networks receptive field. Try changing the downsampling factors to see how it affects the receptive field (eg try combinations of `[1,1]`, `[3,3]`, `[4,4]`, etc)\n",
    "\n",
    "![example_RF](static/example_RF.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d5b6c59-3b2e-4caa-92cc-48162a80908d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# import a UNet class\n",
    "from unet import UNet\n",
    "\n",
    "raw, mask = train_dataset[random.randrange(len(train_dataset))]\n",
    "\n",
    "net_t = raw\n",
    "fovs = []\n",
    "d_factors = [[2,2],[2,2]]\n",
    "\n",
    "# create unet\n",
    "net = UNet(in_channels=1,\n",
    "           num_fmaps=6,\n",
    "           fmap_inc_factors=2,\n",
    "           downsample_factors=d_factors,\n",
    "           padding='same'\n",
    "          )\n",
    "\n",
    "# get unet fovs\n",
    "for level in range(len(d_factors)+1):\n",
    "    fov_tmp, _ = net.rec_fov(level , (1, 1), 1)\n",
    "    fovs.append(fov_tmp[0])\n",
    "\n",
    "fig=plt.figure(figsize=(5, 5))\n",
    "colors = [\"yellow\", \"red\", \"green\"]\n",
    "\n",
    "plt.imshow(np.squeeze(raw), cmap='gray')\n",
    "\n",
    "# visualize receptive field\n",
    "for idx, fov_t in enumerate(fovs):\n",
    "    print(\"Field of view at depth {}: {:3d} (color: {})\".format(idx+1, fov_t, colors[idx]))\n",
    "    xmin = raw.shape[1]/2 - fov_t/2\n",
    "    xmax = raw.shape[1]/2 + fov_t/2\n",
    "    ymin = raw.shape[1]/2 - fov_t/2\n",
    "    ymax = raw.shape[1]/2 + fov_t/2\n",
    "    plt.hlines(ymin, xmin, xmax, color=colors[idx], lw=3)\n",
    "    plt.hlines(ymax, xmin, xmax, color=colors[idx], lw=3)\n",
    "    plt.vlines(xmin, ymin, ymax, color=colors[idx], lw=3)\n",
    "    plt.vlines(xmax, ymin, ymax, color=colors[idx], lw=3)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55352c82-130c-4452-badc-5be3e7ec6a7f",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Set hyperparameters, create model\n",
    "    \n",
    "- Let's start by setting some hyperparameters. Since we are just doing a fg/bg prediction to start, this will be pretty similar to the semantic segmentation exercise. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02950213",
   "metadata": {
    "tags": []
   },
   "source": [
    "<a id='second-todo'></a>\n",
    "<hr style=\"height:2px;\">\n",
    "<div class=\"alert alert-block alert-info\"><h4>Task 1.2</h4>\n",
    "    \n",
    "- Decide how many output channels to have, remember we are starting with a binary task\n",
    "- What loss function and final level activation should we use? Think back to the semantic segmentation exercise.\n",
    "- What type should we ensure our tensors to be? You can see see a list of tensor types [here](https://pytorch.org/docs/stable/tensors.html) - maybe the equivalent of 32-bit floating point :)\n",
    "- For our model, we will create a two level U-Net with the following parameters: (to learn more about the torch layers click [here](https://pytorch.org/docs/stable/nn.html))\n",
    "    - downsample by a factor of 2 in each level\n",
    "    - single input channel\n",
    "    - 32 input feature maps \n",
    "    - multiply by a factor of 2 between levels\n",
    "    - `same` padding (this gives us the same input and output shapes)     \n",
    "    - Since our Unet will have the same number of output features as input features, we need to add a final convolution to get to our desired output feature maps. We should use a final convolution with kernel size of 1 \n",
    "    - To see parameter defs you can run `UNet?`\n",
    "- How many trainable parameters does our network have?\n",
    "\n",
    "Click [here](#first-todo) to go to the previous task.\n",
    "\n",
    "Click [here](#third-todo) to go to the next task.\n",
    "\n",
    "</div>\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78f1cd92",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# set hyperparams\n",
    "\n",
    "out_channels = ... \n",
    "activation = ... \n",
    "loss_fn = ... \n",
    "dtype = ...\n",
    "\n",
    "torch.manual_seed(42)\n",
    "\n",
    "downsample_factors = ... # refer to the init method signature in unet.py module \n",
    "in_channels = ... \n",
    "num_fmaps = ... \n",
    "fmap_inc_factors = ... # refer to the init method signature in unet.py module \n",
    "padding = ... \n",
    "final_kernel_size = ... \n",
    "\n",
    "unet = UNet(\n",
    "        in_channels=in_channels,\n",
    "        num_fmaps=num_fmaps,\n",
    "        fmap_inc_factors=fmap_inc_factors,\n",
    "        downsample_factors=d_factors,\n",
    "        padding=padding)\n",
    "\n",
    "final_conv = torch.nn.Conv2d(\n",
    "    in_channels=num_fmaps,\n",
    "    out_channels=out_channels,\n",
    "    kernel_size=final_kernel_size)\n",
    "\n",
    "net = ... # create your network from your unet and final convolution (hint: torch.nn.Sequential might be useful)\n",
    "\n",
    "net.to(device)\n",
    "\n",
    "summary(net, (in_channels, 64, 64))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b62b36d-00f0-4812-bf59-a94a6b9eebed",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Run training loop\n",
    "    \n",
    "- Here we create a train loop function which calls a train step function on each iteration.\n",
    "    \n",
    "    \n",
    "- As you are familiar with by now, we pass our image or feature into our model to get our logits. We then pass our logits through our final activation in order to get our predictions.\n",
    "    \n",
    "    \n",
    "- Our predictions are the input to our loss, with our target being the ground truth labels. We then backpropagate and step the optimizer. In our case we will use the same train step for validation so we have to make sure to only backprop and step if our net is in train mode. We can control this inside our train loop by setting our net to eval before the validation loop, and back to train once it's done.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41b21fa1-2546-4e8e-a944-1a67b33f2a00",
   "metadata": {},
   "outputs": [],
   "source": [
    "def model_step(model, loss_fn, optimizer, feature, label, activation, prediction_type=None, train_step=True):\n",
    "    \n",
    "    # zero gradients if training\n",
    "    if train_step:\n",
    "        optimizer.zero_grad()\n",
    "    \n",
    "    # forward - pass data through model to get logits\n",
    "    logits = model(feature)\n",
    "    \n",
    "    if prediction_type == \"three_class\":\n",
    "        label=torch.squeeze(label,1)\n",
    "        \n",
    "    # pass logits through final activation to get predictions\n",
    "    predicted = activation(logits)\n",
    "\n",
    "    # pass predictions through loss, compare to ground truth\n",
    "    loss_value = loss_fn(input=predicted, target=label)\n",
    "    \n",
    "    # if training mode, backprop and optimizer step\n",
    "    if train_step:\n",
    "        loss_value.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "    # return outputs and loss\n",
    "    outputs = {\n",
    "        'pred': predicted,\n",
    "        'logits': logits,\n",
    "    }\n",
    "    \n",
    "    return loss_value, outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c7f6008",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(train_loader, val_loader, net, loss_fn, activation, optimizer, dtype, prediction_type=None):\n",
    "\n",
    "    # set train flags, initialize step\n",
    "    net.train() \n",
    "    loss_fn.train()\n",
    "    step = 0\n",
    "\n",
    "\n",
    "    with tqdm(total=training_steps, leave=False, position=0) as pbar:\n",
    "        while step < training_steps:\n",
    "            #reset data loader to get random augmentations\n",
    "            np.random.seed()\n",
    "            tmp_loader = iter(train_loader)\n",
    "            for feature, label in tmp_loader:\n",
    "                label = label.type(dtype)\n",
    "                label = label.to(device)\n",
    "                feature = feature.to(device)\n",
    "                loss_value, pred = model_step(net, loss_fn, optimizer, feature, label, activation, prediction_type)\n",
    "                writer.add_scalar('loss',loss_value.cpu().detach().numpy(),step)\n",
    "                step += 1\n",
    "                pbar.update(1)\n",
    "                if step % 100 == 0:\n",
    "                    net.eval()\n",
    "                    tmp_val_loader = iter(val_loader)\n",
    "                    acc_loss = []\n",
    "                    for feature, label in tmp_val_loader:                    \n",
    "                        label = label.type(dtype)\n",
    "                        label = label.to(device)\n",
    "                        feature = feature.to(device)\n",
    "                        loss_value, _ = model_step(net, loss_fn, optimizer, feature, label, activation, prediction_type, train_step=False)\n",
    "                        acc_loss.append(loss_value.cpu().detach().numpy())\n",
    "                    writer.add_scalar('val_loss',np.mean(acc_loss),step)\n",
    "                    net.train()\n",
    "        print(\"Training complete\")\n",
    "\n",
    "                    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c734f964",
   "metadata": {
    "tags": []
   },
   "source": [
    "<a id='third-todo'></a>\n",
    "<div class=\"alert alert-block alert-info\"><h4>Task 1.3</h4>\n",
    "\n",
    "- Now let's train a model.\n",
    "- Create our data loaders from our datasets. The data loader will take care of batching. We should have 3 dataloaders, one for each dataset we created\n",
    "- The train data loader should use a batch size of 4 (shape = 4, c, h, w) and our val/test data loaders should use a batch size of 1. Set shuffle and pin memory to `True` in the train loader. (for more info on dataloaders see [here](https://pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader))\n",
    "- For now lets just train for 1000 steps\n",
    "- Use a learning rate of 1e-4 and an Adam optimizer\n",
    "- Use the `train` function with all the required parameters to train the model\n",
    "\n",
    "Click [here](#second-todo) to go to the previous task.\n",
    "\n",
    "Click [here](#fourth-todo) to go to the next task.\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff2a58ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_batch_size = ... # set train batch size\n",
    "test_batch_size = ... # set test / val batch size\n",
    "\n",
    "train_loader = ... # create train data loader\n",
    "test_loader = ... # create test data loader\n",
    "val_loader = ... # create val data loader\n",
    "\n",
    "training_steps = ... # set your training setps\n",
    "\n",
    "# create a logdir for each run and a corresponding summary writer\n",
    "logdir = os.path.join(\"logs\", datetime.datetime.now().strftime(\"%Y%m%d-%H%M%S\"))\n",
    "writer = SummaryWriter(logdir)\n",
    "\n",
    "# make sure net and loss are cast to our device (should be gpu, can check by printing device)\n",
    "net = net.to(device)\n",
    "loss_fn = loss_fn.to(device)\n",
    "\n",
    "# set optimizer\n",
    "learning_rate = ... # set your learning rate \n",
    "optimizer = ... # create your optimizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83bb0ce8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# run training loop... (i.e. call train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c5ba248",
   "metadata": {},
   "outputs": [],
   "source": [
    "# To view runs in tensorboard you can call either (uncommented):\n",
    "\n",
    "# %reload_ext tensorboard\n",
    "# %tensorboard --logdir logs\n",
    "\n",
    "# or to view in separate window, run:\n",
    "\n",
    "# !tensorboard --logdir=logs \n",
    "\n",
    "# Note that if running over ssh you will need to also need to add the hostname\n",
    "\n",
    "# !tensorboard --logdir=logs --host hostname"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bba0107c-5a21-421f-8c33-fda6da798048",
   "metadata": {
    "tags": []
   },
   "source": [
    "<h3>Visualize results</h3>\n",
    "    \n",
    "- Here we will run inference using our trained model\n",
    "\n",
    "- We will iterate over the test loader and pass our image through the model\n",
    "\n",
    "- Once we have our prediction (after passing logits through our activation function), we will do a simple thresholding and relabelling to get a segmentation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "daee8d9f-99b6-42fd-87a4-4e3c54c0ee0b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# utility function to relabel connected components\n",
    "from utils import relabel_cc\n",
    "\n",
    "# function to perform otsu thresholding on predictions\n",
    "from skimage.filters import threshold_otsu\n",
    "\n",
    "# make sure net is in eval mode so we don't backprop\n",
    "net.eval()\n",
    "\n",
    "for idx, (image, mask) in enumerate(test_loader):\n",
    "    \n",
    "    # move image to device\n",
    "    image = image.to(device)\n",
    "    \n",
    "    # pass image through network\n",
    "    logits = net(image)\n",
    "    \n",
    "    # pass logits through activation\n",
    "    pred = activation(logits)\n",
    "        \n",
    "    # get our tensors to numpy arrays so we can post-process / visualize\n",
    "    image = image.cpu()\n",
    "    mask = mask.cpu().numpy()\n",
    "    \n",
    "    # we need to detach our predicted tensor\n",
    "    pred = pred.cpu().detach().numpy()\n",
    "\n",
    "    # we also need to remove the batch/channel dimensions from the arrays for visualizing\n",
    "    # (b,c,h,w) -> (h, w)\n",
    "    image = np.squeeze(np.squeeze(image))\n",
    "    mask = np.squeeze(np.squeeze(mask))\n",
    "    pred = np.squeeze(np.squeeze(pred))\n",
    "                \n",
    "    # get threshold value (How does the segmentation change if we change to 0.1 / 0.8?)\n",
    "    thresh = threshold_otsu(pred)\n",
    "    \n",
    "    # fg = prediction greater than / equal to threshold\n",
    "    boundary_mask = pred >= thresh\n",
    "    \n",
    "    # relabel all boundary_mask connected pixels as unique objects\n",
    "    labeled = relabel_cc(boundary_mask)\n",
    "    \n",
    "    # get the corresponding gt labels to visualize\n",
    "    gt_labels = imread(test_loader.dataset.mask_files[idx])\n",
    "    \n",
    "    fig, axes = plt.subplots(1,6,figsize=(20, 20),sharex=True,sharey=True,squeeze=False)\n",
    "        \n",
    "    axes[0][0].imshow(image, cmap='gray')\n",
    "    axes[0][0].title.set_text('Raw')\n",
    "\n",
    "    axes[0][1].imshow(mask)\n",
    "    axes[0][1].title.set_text('GT mask')\n",
    "    \n",
    "    axes[0][2].imshow(create_lut(gt_labels))\n",
    "    axes[0][2].title.set_text('GT seg')\n",
    "    \n",
    "    axes[0][3].imshow(pred)\n",
    "    axes[0][3].title.set_text('Predicted Mask')\n",
    "    \n",
    "    axes[0][4].imshow(boundary_mask)\n",
    "    axes[0][4].title.set_text('Thresholded Mask')\n",
    "\n",
    "    axes[0][5].imshow(create_lut(labeled))\n",
    "    axes[0][5].title.set_text('Predicted Seg')\n",
    "\n",
    "    break"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e027aa06",
   "metadata": {
    "tags": []
   },
   "source": [
    "<hr style=\"height:2px;\"><div class=\"alert alert-block alert-success\"><h2>Checkpoint 1</h2>\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de5f9fbe-33b7-4993-8060-91f2c44e7445",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Part 2: Improving the model\n",
    "\n",
    "- As you can see, our prediction segmentation isn't very good. When objects are tightly packed together, using a simple foreground / background representation gives us results that aren't much better than if we just thresholded our data and relabelled connected components.\n",
    "\n",
    "\n",
    "- So, let's improve our model. We can do a few things to enhance our results:\n",
    "    1. Add more complex representations\n",
    "        * three class\n",
    "        * signed distance transform\n",
    "        * edge affinities\n",
    "    2. Add better augmentations\n",
    "    3. Increase the input size to our network\n",
    "    4. Use a bigger network (eg increase layers, number of feature maps)\n",
    "    5. Train for longer\n",
    "    6. Use a better post-processing strategy (e.g. seeded watershed)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02aa001e-797d-4d6f-9735-b81ad1af2ddc",
   "metadata": {
    "tags": []
   },
   "source": [
    "<h3>Add more augmentations</h3>\n",
    "    \n",
    "- We were using some pretty simple augmentations (crop and flips). Since we want to create a more robust model, we should augment our data more so that it performs better on data that it hasn't seen. This is also a good way to effectively increase our training sample size. \n",
    "- This is a good tutorial for adding useful augmentations: https://albumentations.ai/docs/examples/example_kaggle_salt/\n",
    "- We will still crop and flip our data as before. Additionally, we will:\n",
    "    1. Pad our data if needed (This is useful if our network input size is not compatible with our max pooling layers)\n",
    "    2. Randomly rotate by 90 degrees\n",
    "    3. Transpose\n",
    "    4. Randomly adjust our brightness and contrast\n",
    "- For an intuition about the PadIfNeeded augmentation see the tutorial above. This is for compatibility in our Unet with the max pooling layers. In the following cells we will increase our Unet layers from 2 -> 3. Since we will have 3 max pooling layers, we will need to pad the images by the next closest number that is divisible by 2^3 (8). In our case since we are using a crop size of 64 we won't need to pad since it is divisible by 8. But in the case that our crop size was 65, we would need to take the next number divisible by 8 which would be 72. The PadIfNeeded() augmentation will handle these cases for us, but we need to provide it with the correct padding.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cdc3c6c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# We're using probability of 1 (the padding should always occur if needed)\n",
    "# We're using `border_mode` of 0 to pad it with zeros (reflect is the default which we don't want)\n",
    "\n",
    "def augment_data(raw, mask, padding, crop_size):\n",
    "    \n",
    "    transform = A.Compose([\n",
    "            A.RandomCrop(\n",
    "                width=crop_size,\n",
    "                height=crop_size),\n",
    "            A.PadIfNeeded(\n",
    "                min_height=padding,\n",
    "                min_width=padding,\n",
    "                p=1,\n",
    "                border_mode=0),\n",
    "            A.HorizontalFlip(p=0.3),\n",
    "            A.VerticalFlip(p=0.3),\n",
    "            A.RandomRotate90(p=0.3),\n",
    "            A.Transpose(p=0.3),\n",
    "            A.RandomBrightnessContrast(p=0.3)\n",
    "        ])\n",
    "\n",
    "    transformed = transform(image=raw, mask=mask)\n",
    "\n",
    "    raw, mask = transformed['image'], transformed['mask']\n",
    "    \n",
    "    return raw, mask"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "340d8bc9-825f-4c44-91a7-c08a0206bca9",
   "metadata": {
    "tags": []
   },
   "source": [
    "<h3>Add extra representations</h3>\n",
    "    \n",
    "- Three-class model\n",
    "\n",
    "This is an extension of the basic foreground/background (or two-class) model. In addition a third class is introduced: the boundary. Even if two instances are touching, there is a boundary between them. This way they can be separated. Instead of a single output (where an output of zero is one class and of one is the other class), the network outputs three values, one per class. And the loss function changes from binary to (sparse) categorical cross entropy.\n",
    "    \n",
    "- Signed Distance Transform\n",
    "\n",
    "The label for each pixel is the distance to the closest boundary. The value within instances is negative and outside of instances is positive. As the output is not a probability but an (in principle) unbounded scalar, the mean squared error loss function is used.\n",
    "    \n",
    "- Edge Affinities\n",
    "\n",
    "Here we consider not just the pixel but also its direct neighbors (in 2D the left neighbor and the upper neighbor are sufficient, right and down are redundant with the next pixel's left and upper neighbor). Imagine there is an edge between two pixels if they are in the same class and no edge if not. If we then take all pixels that are directly and indirectly connected by edges, we get an instance. Essentially, we label edges between neighboring pixels as “connected” or “cut”, rather than labeling the pixels themselves. This representation can be useful especially in the case of smaller objects that would otherwise be classified as background pixels. We can use mean squared error. \n",
    "\n",
    "\n",
    "![diff_costs](static/diff_costs.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70f98e43-335a-40c3-8363-6ad6a1d9ca6e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# utility functions to compute signed distance transform and edge affinities\n",
    "from utils import compute_sdt, compute_affinities, erode_border\n",
    "\n",
    "# compute each representation and visualize\n",
    "\n",
    "file = random.choice(train_nuclei)\n",
    "\n",
    "full_mask_nuclei = imread(file)\n",
    "full_raw_nuclei = imread(file.replace('_nuclei_masks', ''))[0]\n",
    "\n",
    "# We're applying a random crop to the image to keep the shape consistent\n",
    "transform = A.Compose([\n",
    "              A.RandomCrop(width=64, height=64),\n",
    "            ])\n",
    "\n",
    "transformed = transform(image=full_raw_nuclei, mask=full_mask_nuclei)\n",
    "          \n",
    "aug_raw, aug_mask = transformed['image'], transformed['mask']\n",
    "\n",
    "# get eroded labels\n",
    "labels, border = erode_border(\n",
    "    aug_mask,\n",
    "    iterations=1,\n",
    "    border_value=1)\n",
    "\n",
    "# get fg/bg classes\n",
    "labels_two_class = (labels != 0)\n",
    "\n",
    "# get border pixels\n",
    "border[border!=0] = 2\n",
    "\n",
    "# combine to get three class \n",
    "labels_three_class = (labels_two_class + border)\n",
    "\n",
    "# compute signed distance transform\n",
    "sdt = compute_sdt(labels)\n",
    "\n",
    "# compute edge affinities\n",
    "affs = compute_affinities(labels, nhood=[[0,1],[1,0]])\n",
    "\n",
    "fig, axes = plt.subplots(1,5,figsize=(20, 10),sharex=True,sharey=True,squeeze=False)\n",
    "\n",
    "for idx, (ds_name, data) in enumerate([\n",
    "    ('raw', aug_raw),\n",
    "    ('fg/bg', labels_two_class),\n",
    "    ('three class', labels_three_class),\n",
    "    ('sdt', sdt),\n",
    "    ('affinities', affs[0] + affs[1])]\n",
    "):\n",
    "\n",
    "    cmap = 'gray' if ds_name == 'raw' else 'viridis'\n",
    "\n",
    "    axes[0][idx].imshow(data.astype(np.float32), cmap=cmap)\n",
    "    axes[0][idx].title.set_text(ds_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec48dc5b",
   "metadata": {
    "tags": []
   },
   "source": [
    "<a id='fourth-todo'></a>\n",
    "<hr style=\"height:2px;\"><div class=\"alert alert-block alert-info\"><h4>Task 2.1</h4>\n",
    "\n",
    "- Add `prediction_type` as a parameter to your dataset\n",
    "- Add a function `create_target` which should take your augmented mask and a prediction type and return the correct representation as float32\n",
    "- Create an extra argument to the dataset `padding_size`. This should be 2^(number of max pooling layers), which will be 2^3 -> `8` for us\n",
    "- Update your augmentation function with the function we defined eatlier\n",
    "- Create a function `get_padding` that takes the `crop_size` and `padding_size` and checks if the crop size is divisible by the padding size. If it is, return the crop size, otherwise get the next higher number that is divisible by the padding size and return this number. \n",
    "- succesful calls should look like:\n",
    "    - `get_padding(64, 8) -> 64`\n",
    "    - `get_padding(65, 8) -> 72`\n",
    "    - `get_padding(72, 8) -> 72`\n",
    "    - `get_padding(73, 8) -> 80`\n",
    "- The returned padding should then be passed in as the `min_height` and `min_width` for the augmentation\n",
    "- Your `__getitem__` should call `create_target` after augmenting your data\n",
    "- Try each prediction type \n",
    "- What happens if you increase crop size to 65? - how does the padding come into play?\n",
    "\n",
    "Click [here](#third-todo) if you need to go back to the previous task.\n",
    "\n",
    "Click [here](#fifth-todo) if you need to go to the next task.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20072fbd",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TissueNetDataset(Dataset):\n",
    "    def __init__(self,\n",
    "                 root_dir,\n",
    "                 split='train',\n",
    "                 mask='nuclei',\n",
    "                 crop_size=None,\n",
    "                 val_split=False,\n",
    "                 prediction_type='two_class',\n",
    "                 padding_size=8\n",
    "                ):\n",
    "        \n",
    "        self.split = split\n",
    "        self.crop_size = crop_size\n",
    "        \n",
    "        self.mask_files = natsorted(glob(os.path.join(root_dir, split, f'*{mask}*')))\n",
    "        self.raw_files = [i for i in natsorted(glob(os.path.join(root_dir, split, '*.tif'))) if 's.tif' not in i]\n",
    "        \n",
    "        # Add a parameter `prediction_type` to define which representation to choose.\n",
    "        # Default should be `two_class`\n",
    "        \n",
    "        # Add a parameter `padding_size` to define when to pad the image\n",
    "        # Default should be 8 (2^3 for 3 max pooling layers)\n",
    "\n",
    "        if split == 'test':\n",
    "            if val_split:\n",
    "                self.mask_files = self.mask_files[:10]\n",
    "                self.raw_files = self.raw_files[:10]\n",
    "        else:\n",
    "            self.mask_files = self.mask_files[10:]\n",
    "            self.raw_files = self.raw_files[10:]\n",
    "        \n",
    "        \n",
    "    def __len__(self):\n",
    "        return len(self.raw_files)\n",
    "    \n",
    "    \n",
    "    def get_padding(self, crop_size, padding_size):\n",
    "    \n",
    "        padding = ... # calculate your padding\n",
    "        \n",
    "        # hint:\n",
    "        \n",
    "        # 1. get padding quotient (crop size / padding size)\n",
    "        # 2. if crop size IS NOT evenly divisible by padding size, our padding\n",
    "        #    is equal to (padding size * (padding quotient + 1))\n",
    "        # 3. if crop size IS evenly divisible by padding size, out padding is\n",
    "        #    is equal to our crop size\n",
    "    \n",
    "        return padding\n",
    "\n",
    "    \n",
    "    def create_target(self, mask, prediction_type):\n",
    "        \n",
    "        mask, border = ... # erode your labels, return inner and border\n",
    "        \n",
    "        if prediction_type == 'two_class':\n",
    "            mask = ... # get two class\n",
    "        elif prediction_type == 'three_class':\n",
    "            mask = ... # get three class\n",
    "        elif prediction_type == 'sdt':\n",
    "            mask = ... # get sdt using the function defined above\n",
    "        elif prediction_type == 'affs':\n",
    "            mask = ... # get affs using the function defined above\n",
    "            \n",
    "        mask = ... # cast mask to float32\n",
    "        \n",
    "        return mask  \n",
    "    \n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        raw_file = ... # get raw file at index\n",
    "        mask_file = ... # get mask file at index\n",
    "        \n",
    "        raw = imread(raw_file)\n",
    "        mask = imread(mask_file)\n",
    "        \n",
    "        raw = ... # get nuclei channel, dont forget to cast to float32\n",
    "        \n",
    "        if self.split == 'train':\n",
    "            padding = ... # get padding using the function we defined above \n",
    "            raw, mask = ... # get your augment your data using the correct parameters\n",
    "        \n",
    "        mask = ... # create your learning representation\n",
    "        \n",
    "        raw = ... # add raw channel dimension\n",
    "        mask = ... # add mask channel dimension (**if not** using affs since you get two channels anyway)\n",
    "            \n",
    "        return raw, mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aeced379-98f3-4d32-b4ba-c59cf579a748",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "prediction_type = ... # try each of prediction types we defined in create_target function\n",
    "crop_size = ... # try some different crop sizes, how does the crop size affect the padding? \n",
    "\n",
    "train_dataset = TissueNetDataset(\n",
    "    root_dir='woodshole',\n",
    "    split='train',\n",
    "    crop_size=crop_size,\n",
    "    prediction_type=prediction_type)\n",
    "\n",
    "raw, mask = train_dataset[random.randrange(len(train_dataset))]\n",
    "\n",
    "fig, axes = plt.subplots(1,2,figsize=(10, 10),sharex=True,sharey=True,squeeze=False)\n",
    "\n",
    "axes[0][0].imshow(np.squeeze(raw), cmap='gray')\n",
    "axes[0][0].title.set_text('Raw')\n",
    "\n",
    "if mask.shape[0] == 1:\n",
    "    axes[0][1].imshow(np.squeeze(mask))\n",
    "    axes[0][1].title.set_text('Mask')\n",
    "else:\n",
    "    # affs has two channels (x/y)\n",
    "    axes[0][1].imshow(mask[0]+mask[1])\n",
    "    axes[0][1].title.set_text('Mask')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e33002a1-0937-4f89-a0c4-33de721960ae",
   "metadata": {
    "tags": []
   },
   "source": [
    "<h3> Add extra hyperparameters based on prediction type</h3>\n",
    "    \n",
    "- Now that we have more representations, we need to be sure that our hyperparams are consistent with whichever representation we choose. \n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "357d9786",
   "metadata": {
    "tags": []
   },
   "source": [
    "<a id='fifth-todo'></a>\n",
    "<hr style=\"height:2px;\"><div class=\"alert alert-block alert-info\"><h4>Task 2.2</h4>\n",
    "\n",
    "- Create a function `get_hyperparams` that takes your `prediction_type` and returns a dictionary mapping our hyperparameters:   \n",
    "    ```\n",
    "    params = {\n",
    "        'out_channels': out_channels,\n",
    "        'activation': activation,\n",
    "        'loss_fn', loss_fn,\n",
    "        'dtype': dtype\n",
    "     }\n",
    "     ```\n",
    "- You already know what these need to be for the two class representation, let's figure out the others\n",
    "- **Hint**: We will choose from the following loss functions: `nn.CrossEntropyLoss` and `nn.MSELoss`\n",
    "- **Hint**: We will choose from the following final activations: `tanh`, `softmax`, `sigmoid`\n",
    "- **Three class**:\n",
    "    - We are now doing multi class classification, what type of loss function should we use?\n",
    "    - Our number of out channels should be equal to the number of classes we are trying to predict\n",
    "    - What final activation should we use? What dimension should it be computed on? Remember our tensors will have shape (b,c,h,w) or (0,1,2,3)\n",
    "    - What dtype should we have? It should be the torch equivalent of a 64-bit integer (signed) (see [here](https://pytorch.org/docs/stable/tensors.html) if stuck with tensor types)   \n",
    "- **Sdt**:\n",
    "    - Since we are computing a **signed** distance transform, we want to have negative values outside of our objects. Therefore we want an activation that can give us values between -1 and 1. \n",
    "    - We are doing regression now, what loss function can we use?\n",
    "    - Our output will be a single number per pixel, so how many output channels should we have?\n",
    "    - Our dtype should be the torch equivalent of a 32-bit floating point   \n",
    "- **Affs**:\n",
    "    - We are doing regression again, what should our loss be?\n",
    "    - Our outputs will be between 0 and 1, what activation should we use?\n",
    "    - We can use the same dtype as Sdt\n",
    "    - We will have both x and y affinities, so how many channels should we have?\n",
    " \n",
    "\n",
    "Click [here](#fourth-todo) to go to the previous task.\n",
    "\n",
    "Click [here](#sixth-todo) to go to the next task.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a280ceb",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_hyperparams(prediction_type):\n",
    "    \n",
    "    if prediction_type == \"two_class\":\n",
    "        out_channels = 1\n",
    "        activation = torch.nn.Sigmoid()\n",
    "        loss_fn = torch.nn.BCELoss()\n",
    "        dtype = torch.FloatTensor\n",
    "        \n",
    "    elif prediction_type == \"three_class\":\n",
    "        ... # get params\n",
    "        \n",
    "    elif prediction_type == \"sdt\":\n",
    "        ... # get params\n",
    "        \n",
    "    elif prediction_type == \"affs\":\n",
    "        ... # get params\n",
    "        \n",
    "    else:\n",
    "        raise ValueError('Choose from one of the following prediction types: two_class, three_class, sdt, affs')\n",
    "        \n",
    "    params = ... # get dict\n",
    "    \n",
    "    return params"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71a7ea72-0a89-4900-b736-1ccdece8bb82",
   "metadata": {},
   "outputs": [],
   "source": [
    "prediction_type = 'affs'\n",
    "\n",
    "params = get_hyperparams(prediction_type)\n",
    "\n",
    "print(params)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3533b935-dea0-4115-91d6-d273362db9ad",
   "metadata": {
    "tags": []
   },
   "source": [
    "<h3>Increase patch crop size</h3>\n",
    "    \n",
    "- Before we were using a smaller patch crop size (64). Since we are training a 2d network with a relatively small batch number (4), it is not such a big deal to increase our crop size (128) to let our network see more data.\n",
    "- **Note**, that increasing the crop size won't affect the receptive field of the network.\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "837093e4-5975-4821-a26e-59fe141051c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# try each prediction type\n",
    "prediction_type = 'affs'\n",
    "crop_size = 128\n",
    "\n",
    "train_dataset = TissueNetDataset(root_dir='woodshole', split='train', crop_size=crop_size, prediction_type=prediction_type)\n",
    "\n",
    "raw, mask = train_dataset[random.randrange(len(train_dataset))]\n",
    "\n",
    "fig, axes = plt.subplots(1,2,figsize=(10, 10),sharex=True,sharey=True,squeeze=False)\n",
    "\n",
    "axes[0][0].imshow(np.squeeze(raw), cmap='gray')\n",
    "axes[0][0].title.set_text('Raw')\n",
    "\n",
    "if mask.shape[0] == 1:\n",
    "    axes[0][1].imshow(np.squeeze(mask))\n",
    "    axes[0][1].title.set_text('Mask')\n",
    "else:\n",
    "    # affs has two channels (x/y)\n",
    "    axes[0][1].imshow(mask[0]+mask[1])\n",
    "    axes[0][1].title.set_text('Mask')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f795ce0-440f-4c27-ba5a-3b0638b3abca",
   "metadata": {
    "tags": []
   },
   "source": [
    "<h3>Increase features in network and train longer</h3>\n",
    "    \n",
    "- Before we were training with a pretty small network, eg two downsampling levels. Let's see how the receptive fields change as we increase our network to 3 levels\n",
    "    \n",
    "    \n",
    "- Try changing the downsampling factors as before. How does it change the receptive field?\n",
    "\n",
    "![deep_RF](static/deep_RF.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0598e3d6-0c55-4753-a9b8-7b5cd3bd7dc0",
   "metadata": {},
   "outputs": [],
   "source": [
    "raw, mask = train_dataset[random.randrange(len(train_dataset))]\n",
    "\n",
    "net_t = raw\n",
    "fovs = []\n",
    "d_factors = [[2,2],[2,2],[2,2]]\n",
    "\n",
    "net = UNet(in_channels=1,\n",
    "           num_fmaps=6,\n",
    "           fmap_inc_factors=2,\n",
    "           downsample_factors=d_factors,\n",
    "           padding='same'\n",
    "          )\n",
    "\n",
    "for level in range(len(d_factors)+1):\n",
    "    fov_tmp, _ = net.rec_fov(level , (1, 1), 1)\n",
    "    fovs.append(fov_tmp[0])\n",
    "\n",
    "fig=plt.figure(figsize=(8, 8))\n",
    "colors = [\"yellow\", \"red\", \"green\", \"blue\"]\n",
    "\n",
    "plt.imshow(np.squeeze(raw), cmap='gray')\n",
    "\n",
    "for idx, fov_t in enumerate(fovs):\n",
    "    print(\"Field of view at depth {}: {:3d} (color: {})\".format(idx+1, fov_t, colors[idx]))\n",
    "    xmin = raw.shape[1]/2 - fov_t/2\n",
    "    xmax = raw.shape[1]/2 + fov_t/2\n",
    "    ymin = raw.shape[1]/2 - fov_t/2\n",
    "    ymax = raw.shape[1]/2 + fov_t/2\n",
    "    plt.hlines(ymin, xmin, xmax, color=colors[idx], lw=3)\n",
    "    plt.hlines(ymax, xmin, xmax, color=colors[idx], lw=3)\n",
    "    plt.vlines(xmin, ymin, ymax, color=colors[idx], lw=3)\n",
    "    plt.vlines(xmax, ymin, ymax, color=colors[idx], lw=3)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a306ac4d",
   "metadata": {
    "tags": []
   },
   "source": [
    "<a id='sixth-todo'></a>\n",
    "\n",
    "<div class=\"alert alert-block alert-info\"><h4>Task 2.3</h4>\n",
    "    \n",
    "- Choose a prediction type (other than two_class) to use and get the corresponding parameters\n",
    "- Add another layer to your network with the same downsampling factors\n",
    "- Increase your multiplication factor between layers to 3\n",
    "- Make sure you correctly set the output channels in your final convolution (hint: get it from your param dict)\n",
    "- How many trainable parameters do we have now? How does this compare to when we used two layers and a mult factor of 2 instead of 3?\n",
    "- Create your datasets and loaders as before. Increase crop patch size (eg 64 -> 128)    \n",
    "- Use the same learning rate and optimizer as before   \n",
    "- Train for longer (eg 1000 -> 3000 steps)\n",
    "\n",
    "Click [here](#fifth-todo) to go to the previous task.\n",
    "\n",
    "Click [here](#seventh-todo) to go to the next task.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc6fced5",
   "metadata": {},
   "outputs": [],
   "source": [
    "prediction_type = ... \n",
    "params = get_hyperparams(prediction_type)\n",
    "\n",
    "torch.manual_seed(42)\n",
    "\n",
    "downsample_factors = ... # refer to the init method signature in unet.py module \n",
    "in_channels=1\n",
    "num_fmaps=32\n",
    "fmap_inc_factors= ... # refer to the init method signature in unet.py module \n",
    "padding='same'\n",
    "final_kernel_size=1\n",
    "out_channels = params['out_channels']\n",
    "\n",
    "unet = UNet(\n",
    "        in_channels=in_channels,\n",
    "        num_fmaps=num_fmaps,\n",
    "        fmap_inc_factors=fmap_inc_factors,\n",
    "        downsample_factors=d_factors,\n",
    "        padding=padding)\n",
    "\n",
    "final_conv = torch.nn.Conv2d(\n",
    "    in_channels=num_fmaps,\n",
    "    out_channels=out_channels,\n",
    "    kernel_size=final_kernel_size)\n",
    "\n",
    "net = torch.nn.Sequential(unet, final_conv)\n",
    "\n",
    "net = net.to(device)\n",
    "\n",
    "summary(net, (in_channels, 128, 128))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74f2bf28",
   "metadata": {},
   "outputs": [],
   "source": [
    "training_steps = ...\n",
    "logdir = os.path.join(\"logs\", datetime.datetime.now().strftime(\"%Y%m%d-%H%M%S\"))\n",
    "writer = SummaryWriter(logdir)\n",
    "\n",
    "net = net.to(device)\n",
    "loss_fn = params['loss_function'].to(device)\n",
    "activation = params['activation']\n",
    "dtype = params['dtype']\n",
    "\n",
    "crop_size = ...\n",
    "\n",
    "### create datasets\n",
    "train_dataset = TissueNetDataset(root_dir='woodshole', split='train', crop_size=crop_size, prediction_type=prediction_type)\n",
    "test_dataset = TissueNetDataset(root_dir='woodshole', split='test', prediction_type=prediction_type)\n",
    "val_dataset = TissueNetDataset(root_dir='woodshole', split='test', val_split=True, prediction_type=prediction_type)\n",
    "\n",
    "batch_size = 4\n",
    "\n",
    "# make dataloaders\n",
    "train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, pin_memory=True)\n",
    "test_loader = DataLoader(test_dataset, batch_size=1)\n",
    "val_loader = DataLoader(val_dataset, batch_size=1)\n",
    "\n",
    "# set optimizer\n",
    "learning_rate = 1e-4\n",
    "optimizer = torch.optim.Adam(net.parameters(), lr=learning_rate)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18fb5f75",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# run training loop\n",
    "train(train_loader, val_loader, net, loss_fn, activation, optimizer, dtype, prediction_type)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d302c119-a72b-472e-813e-70e12a288aa6",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Visualize predictions\n",
    "\n",
    "net.eval()\n",
    "\n",
    "for idx, (image, mask) in enumerate(test_loader):\n",
    "    image = image.to(device)\n",
    "    logits = net(image)\n",
    "    pred = activation(logits)\n",
    "        \n",
    "    image = np.squeeze(image.cpu())\n",
    "    mask = np.squeeze(mask.cpu().numpy())\n",
    "    \n",
    "    pred = np.squeeze(pred.cpu().detach().numpy())\n",
    "        \n",
    "    fig, axes = plt.subplots(1,3,figsize=(20, 20),sharex=True,sharey=True,squeeze=False)\n",
    "    \n",
    "    axes[0][0].imshow(image, cmap='gray')\n",
    "    axes[0][0].title.set_text('Raw')\n",
    "    \n",
    "    if prediction_type == 'three_class':\n",
    "        # get indices of the maximum values along channel axis\n",
    "        pred = np.argmax(pred, axis=0)\n",
    "        \n",
    "    if prediction_type == 'affs':\n",
    "        axes[0][1].imshow(mask[0] + mask[1])\n",
    "        axes[0][1].title.set_text('GT mask')\n",
    "        axes[0][2].imshow(pred[0] + pred[1])\n",
    "        axes[0][2].title.set_text('Predicted')\n",
    "        \n",
    "    else:\n",
    "        axes[0][1].imshow(mask)\n",
    "        axes[0][1].title.set_text('GT mask')\n",
    "        axes[0][2].imshow(pred)\n",
    "        axes[0][2].title.set_text('Predicted')\n",
    "      \n",
    "    if idx == 2:\n",
    "        break"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "51aad2aa",
   "metadata": {
    "tags": []
   },
   "source": [
    "<hr style=\"height:2px;\"><div class=\"alert alert-block alert-success\"><h2>Checkpoint 2</h2>\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3ffca7a-7dce-476e-99f2-c351b02aaade",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Part 3: Post-processing / further improvements\n",
    "\n",
    "- Before we were just thresholding and relabeling connected components. Now we will also see a more advanced post-processing strategy called watershed which works with affs.\n",
    "\n",
    "\n",
    "- We also want to gauge our model performance so we will introduce some evaluation methods for instance segmentation.\n",
    "\n",
    "\n",
    "- Finally, up until now we were just using a single input channel to our network - but since we have multiple channels in our raw data we should leverage them. You will get the opportunity to put everything together to try to improve your model."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e6665f2-1a3b-42d5-ac8c-27183cb5abce",
   "metadata": {
    "tags": []
   },
   "source": [
    "<h3> Watershed</h3>\n",
    "    \n",
    "- Before we were just thresholding our predictions and then relabeling connected components. This is a totally fine approach in the cases where we don't have touching objects. Now we will use a better approach commonly used for instance segmentation called seeded watershed. See here for a nice overview: https://scikit-image.org/docs/stable/auto_examples/segmentation/plot_watershed.html\n",
    "    \n",
    "    \n",
    "- To compute our seeded watershed, we first need to get a boundary mask from our predictions. This is done slightly differently for the various representations, but generally speaking our boundary mask will just be a boolean indicating our foreground regions. From this boundary mask we compute boundary distances using a distance transform. These will then give us local maxima that can be used to extract seed points. The watershed algorithm then expands each seed out in a local \"basin\" until the segments touch.\n",
    "    \n",
    "    \n",
    "- Because of this, it is often not sufficient to use watershed alone on complex datasets. In most cases the resulting objects are referred to as fragments (or supervoxels), which can then be stitched together using the underlying predictions as edge weights through a process called agglomeration.\n",
    "    \n",
    "    \n",
    "- Agglomeration is out of the scope of this exercise, but you can find a nice overview here: https://scikit-image.org/docs/stable/auto_examples/segmentation/plot_boundary_merge.html. A good challenge for future learning is to implement your own agglomeration using predictions and supervoxels. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "692be5ec-8f3f-4ddb-9e28-02a554dfb023",
   "metadata": {},
   "outputs": [],
   "source": [
    "# utility functions to get boundary mask, and seg using watershed\n",
    "from utils import get_boundary_mask, watershed_from_boundary_distance\n",
    "\n",
    "# function to compute exact euclidean distance transform\n",
    "from scipy.ndimage import distance_transform_edt\n",
    "\n",
    "# get segmentations\n",
    "\n",
    "net.eval()\n",
    "\n",
    "for idx, (image, mask) in enumerate(test_loader):\n",
    "    image = image.to(device)\n",
    "    logits = net(image)\n",
    "    pred = activation(logits)\n",
    "        \n",
    "    image = np.squeeze(image.cpu())\n",
    "    mask = np.squeeze(mask.cpu().numpy())\n",
    "        \n",
    "    pred = np.squeeze(pred.cpu().detach().numpy())\n",
    "    \n",
    "    # feel free to try different thresholds\n",
    "    thresh = np.mean(pred)\n",
    "    \n",
    "    # get boundary mask\n",
    "    boundary_mask = get_boundary_mask(pred, prediction_type, thresh=thresh)\n",
    "    \n",
    "    # get segmentation   \n",
    "    if prediction_type == \"affs\":\n",
    "        # get boundary distances\n",
    "        boundary_distances = distance_transform_edt(boundary_mask)\n",
    "        \n",
    "        seg = watershed_from_boundary_distance(\n",
    "            boundary_distances,\n",
    "            boundary_mask\n",
    "        )\n",
    "    else:\n",
    "        seg = relabel_cc(boundary_mask)\n",
    "    \n",
    "    gt_labels = imread(test_loader.dataset.mask_files[idx])\n",
    "    \n",
    "    fig, axes = plt.subplots(1,3,figsize=(20, 20),sharex=True,sharey=True,squeeze=False)\n",
    "    \n",
    "    axes[0][0].imshow(image, cmap='gray')\n",
    "    axes[0][0].title.set_text('Raw')\n",
    "    \n",
    "    axes[0][1].imshow(create_lut(gt_labels))\n",
    "    axes[0][1].title.set_text('GT Labels')\n",
    "    \n",
    "    axes[0][2].imshow(create_lut(seg))\n",
    "    axes[0][2].title.set_text('Predicted Labels')\n",
    "    \n",
    "    if idx == 2:\n",
    "        break"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f8e17356-b7df-4cb3-98ad-5231f0b915f8",
   "metadata": {
    "tags": []
   },
   "source": [
    "<h3>Evaluate</h3>\n",
    "    \n",
    "- There are several ways to evaluate accuracy of models for instance segmentation.\n",
    "    \n",
    "    \n",
    "- For our purposes, we will calculate the intersection over union (IoU) between the ground truth labels and the predicted labels. Here is a nice overview: https://www.jeremyjordan.me/evaluating-image-segmentation-models/\n",
    "    \n",
    "    \n",
    "- We then choose a threshold on the IoU for defining a match, e.g. if we choose a threshold of 0.5, then an IoU >= 0.5 means a true positive for a ground truth label and a predicted label.\n",
    "    \n",
    "    \n",
    "- Using IoU and the threshold, we match predicted labels to ground truth labels, and then evaluate:\n",
    "    1. True positives\n",
    "    2. False positives\n",
    "    3. False negatives\n",
    "    4. Precision\n",
    "    5. Recall\n",
    "    6. Average precision\n",
    "    \n",
    "    \n",
    "- These will already give a good indication of model performance. You can easily look up some information on each of these metrics, and you will have a good idea about why they are used following the previous exercises and lectures.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0933b0ef-6731-4fc6-a746-21f45acde9d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# import utility function to evaluate several metrics\n",
    "from utils import evaluate\n",
    "\n",
    "# Evaluate on a single batch\n",
    "\n",
    "net.eval()\n",
    "\n",
    "for idx, (image, mask) in enumerate(test_loader):\n",
    "    image = image.to(device)\n",
    "    logits = net(image)\n",
    "    pred = activation(logits)\n",
    "        \n",
    "    image = np.squeeze(image.cpu())\n",
    "    mask = np.squeeze(mask.cpu().numpy())\n",
    "        \n",
    "    pred = np.squeeze(pred.cpu().detach().numpy())\n",
    "        \n",
    "    thresh = np.mean(pred)\n",
    "                  \n",
    "    boundary_mask = get_boundary_mask(pred, prediction_type, thresh=thresh)\n",
    "    \n",
    "    if prediction_type == \"affs\":\n",
    "        boundary_distances = distance_transform_edt(boundary_mask)\n",
    "\n",
    "        pred_labels = watershed_from_boundary_distance(\n",
    "            boundary_distances,\n",
    "            boundary_mask\n",
    "        )\n",
    "        \n",
    "    else:\n",
    "        pred_labels = relabel_cc(boundary_mask)\n",
    "    \n",
    "    gt_labels = imread(test_loader.dataset.mask_files[idx])\n",
    "    \n",
    "    ap, precision, recall, tp, fp, fn = evaluate(gt_labels, pred_labels)\n",
    "    \n",
    "    print(\n",
    "        f'Computed with IoU threshold = 0.5\\n',\n",
    "        f'Average precision: {ap} \\n',\n",
    "        f'Precision: {precision} \\n',\n",
    "        f'Recall: {recall} \\n',\n",
    "        f'True positives: {tp} \\n',\n",
    "        f'False positives: {fp} \\n',\n",
    "        f'False negatives: {fn} \\n'\n",
    "    )\n",
    "        \n",
    "    fig, axes = plt.subplots(1,3,figsize=(20, 20),sharex=True,sharey=True,squeeze=False)\n",
    "    \n",
    "    axes[0][0].imshow(image, cmap='gray')\n",
    "    axes[0][0].title.set_text('Raw')\n",
    "    \n",
    "    axes[0][1].imshow(create_lut(gt_labels))\n",
    "    axes[0][1].title.set_text('GT Labels')\n",
    "    \n",
    "    axes[0][2].imshow(create_lut(pred_labels))\n",
    "    axes[0][2].title.set_text('Predicted Labels')\n",
    "    \n",
    "    break"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ff4f94a-c301-4007-864a-8bc1b494b7cd",
   "metadata": {
    "tags": []
   },
   "source": [
    "<h3>Loop over batches</h3>\n",
    "    \n",
    "- Now we can loop over all test set images and get the the average model precision\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5ae3f5c-9980-4eaf-a54f-ff4da92dfde7",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "avg = 0.0\n",
    "\n",
    "for idx, (image, mask) in enumerate(test_loader):\n",
    "    image = image.to(device)\n",
    "    logits = net(image)\n",
    "    pred = activation(logits)\n",
    "        \n",
    "    image = np.squeeze(image.cpu())\n",
    "    mask = np.squeeze(mask.cpu().numpy())\n",
    "        \n",
    "    pred = np.squeeze(pred.cpu().detach().numpy())\n",
    "        \n",
    "    thresh = np.mean(thresh)\n",
    "            \n",
    "    boundary_mask = get_boundary_mask(pred, prediction_type, thresh)\n",
    "    \n",
    "    if prediction_type == \"affs\":\n",
    "        boundary_distances = distance_transform_edt(boundary_mask)\n",
    "\n",
    "        pred_labels = watershed_from_boundary_distance(\n",
    "            boundary_distances,\n",
    "            boundary_mask\n",
    "        )\n",
    "        \n",
    "    else:\n",
    "        pred_labels = relabel_cc(boundary_mask)\n",
    "    \n",
    "    gt_labels = imread(test_loader.dataset.mask_files[idx])\n",
    "    \n",
    "    ap, precision, recall, tp, fp, fn = evaluate(gt_labels, pred_labels)\n",
    "    \n",
    "    avg += ap\n",
    "    \n",
    "    print(\n",
    "        f'Average precision: {ap} \\n',\n",
    "        f'Precision: {precision} \\n',\n",
    "        f'Recall: {recall} \\n',\n",
    "        f'True positives: {tp} \\n',\n",
    "        f'False positives: {fp} \\n',\n",
    "        f'False negatives: {fn} \\n'\n",
    "    )\n",
    "        \n",
    "    fig, axes = plt.subplots(1,3,figsize=(20, 20),sharex=True,sharey=True,squeeze=False)\n",
    "    \n",
    "    axes[0][0].imshow(image, cmap='gray')\n",
    "    axes[0][0].title.set_text('Raw')\n",
    "    \n",
    "    axes[0][1].imshow(create_lut(gt_labels))\n",
    "    axes[0][1].title.set_text('GT Labels')\n",
    "    \n",
    "    axes[0][2].imshow(create_lut(pred_labels))\n",
    "    axes[0][2].title.set_text('Predicted Labels')\n",
    "    \n",
    "    plt.show()\n",
    "        \n",
    "avg /= (idx+1)\n",
    "    \n",
    "print(\"average precision on test set: {}\".format(avg))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f4389b9",
   "metadata": {
    "tags": []
   },
   "source": [
    "<h3>Use both channels of raw data</h3>\n",
    "\n",
    "- Up until now, we have only been using the nuclei channel of the raw data as input into our network. But if we have extra channels available, it will help our network to see them. We should give our network as much information as possible to learn from, even if it is only tasked with learning a single channel output.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c2c19d7",
   "metadata": {
    "tags": []
   },
   "source": [
    "<a id='seventh-todo'></a>\n",
    "\n",
    "<div class=\"alert alert-block alert-info\"><h4>Task 3.1</h4>\n",
    "\n",
    "- Update your dataset to use both channels of raw data\n",
    "- It is very important to keep track of the shape of your data as it passes through.\n",
    "- By default our raw data is (C, H, W) but albumentations expects (H,W,C) for rgb data\n",
    "- Our mask data is H,W by default so we need to add a channel dimension\n",
    "- Following augmentation, we then need to get both our raw and mask back back to (C,H,W) for training\n",
    "- Make sure when creating your target representation that you handle the mask channel correctly (eg pass `mask[0]` instead of `mask`)\n",
    "\n",
    "Click [here](#sixth-todo) to go to the previous task.\n",
    "\n",
    "Click [here](#eighth-todo) to go to the next task.\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f701c49",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the dataset class \n",
    "\n",
    "class TissueNetDataset(Dataset):\n",
    "    def __init__(self,\n",
    "                 root_dir,\n",
    "                 split='train',\n",
    "                 mask='nuclei',\n",
    "                 crop_size=None,\n",
    "                 val_split=False,\n",
    "                 prediction_type='two_class',\n",
    "                 padding_size=8\n",
    "                ):\n",
    "        \n",
    "        self.split = split\n",
    "        self.mask_files = natsorted(glob(os.path.join(root_dir, split, f'*{mask}*')))\n",
    "        self.raw_files = [i for i in natsorted(glob(os.path.join(root_dir, split, '*.tif'))) if 's.tif' not in i]\n",
    "        self.crop_size = crop_size\n",
    "        self.prediction_type = prediction_type\n",
    "        self.padding_size = padding_size\n",
    "        \n",
    "        if split == 'test':\n",
    "            if val_split:\n",
    "                self.mask_files = self.mask_files[:10]\n",
    "                self.raw_files = self.raw_files[:10]\n",
    "            else:\n",
    "                self.mask_files = self.mask_files[10:]\n",
    "                self.raw_files = self.raw_files[10:]\n",
    "\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.raw_files)\n",
    "    \n",
    "    \n",
    "    def get_padding(self, crop_size, padding_size):\n",
    "    \n",
    "        # quotient\n",
    "        q = int(crop_size / padding_size)\n",
    "    \n",
    "        if crop_size % padding_size != 0:\n",
    "            padding = (padding_size * (q + 1))\n",
    "        else:\n",
    "            padding = crop_size\n",
    "    \n",
    "        return padding\n",
    "\n",
    "\n",
    "    def create_target(self, mask, prediction_type):\n",
    "        \n",
    "        mask, border = erode_border(\n",
    "                    mask,\n",
    "                    iterations=1,\n",
    "                    border_value=1)\n",
    "        \n",
    "        if self.prediction_type == 'two_class':\n",
    "            mask = (mask != 0)\n",
    "\n",
    "        elif self.prediction_type == 'three_class':\n",
    "            labels_two_class = (mask != 0)\n",
    "            border[border!=0] = 2\n",
    "            \n",
    "            mask = labels_two_class + border\n",
    "\n",
    "        elif self.prediction_type == 'sdt':\n",
    "            mask = compute_sdt(mask)\n",
    "\n",
    "        elif self.prediction_type == 'affs':\n",
    "            mask = compute_affinities(mask, nhood=[[0,1],[1,0]])\n",
    "\n",
    "        else:\n",
    "            raise Exception('Choose from one of the following prediction types: two_class, three_class, sdt, affs')\n",
    "        \n",
    "        return mask.astype(np.float32)\n",
    "    \n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        raw_file = self.raw_files[idx]\n",
    "        mask_file = self.mask_files[idx]\n",
    "        \n",
    "        raw = imread(raw_file)\n",
    "        mask = imread(mask_file)\n",
    "        \n",
    "        raw = ... # make dimensions are in the right order, transpose could be useful. \n",
    "        # And don't forget to cast to float32 \n",
    "        mask = np.expand_dims(mask, axis=-1)\n",
    "        \n",
    "        if self.split == 'train':\n",
    "            padding = self.get_padding(self.crop_size, self.padding_size)\n",
    "            raw, mask = augment_data(raw, mask, padding, self.crop_size)\n",
    "            \n",
    "        raw = ... # yet again, make sure the dimensions are in the right order\n",
    "        mask = ... # yet again, make sure the dimensions are in the right order\n",
    "        \n",
    "        mask = self.create_target(mask[0], self.prediction_type)\n",
    "        \n",
    "        if self.prediction_type != 'affs':\n",
    "            mask = np.expand_dims(mask, axis=0)\n",
    "                                        \n",
    "        return raw, mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33fc4037",
   "metadata": {},
   "outputs": [],
   "source": [
    "# try each prediction type\n",
    "prediction_type = 'affs'\n",
    "crop_size = 128\n",
    "\n",
    "train_dataset = TissueNetDataset(root_dir='woodshole', split='train', crop_size=crop_size, prediction_type=prediction_type)\n",
    "\n",
    "raw, mask = train_dataset[random.randrange(len(train_dataset))]\n",
    "\n",
    "raw = np.vstack((raw, np.zeros_like(raw)[:1]))\n",
    "raw = raw.transpose(1,2,0)\n",
    "\n",
    "fig, axes = plt.subplots(1,2,figsize=(10, 10),sharex=True,sharey=True,squeeze=False)\n",
    "\n",
    "axes[0][0].imshow(np.squeeze(raw), cmap='gray')\n",
    "axes[0][0].title.set_text('Raw')\n",
    "\n",
    "if mask.shape[0] == 1:\n",
    "    axes[0][1].imshow(np.squeeze(mask))\n",
    "    axes[0][1].title.set_text('Mask')\n",
    "else:\n",
    "    # affs has two channels (x/y)\n",
    "    axes[0][1].imshow(mask[0]+mask[1])\n",
    "    axes[0][1].title.set_text('Mask')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c8dd0bee",
   "metadata": {
    "tags": []
   },
   "source": [
    "<a id='eighth-todo'></a>\n",
    "\n",
    "<div class=\"alert alert-block alert-info\"><h4>Task 3.2</h4>\n",
    "    \n",
    "- Create your network, hyperparameters and data loaders as last time (eg 3 levels, 3 mult factor, 128 crop size, 3k iterations, etc). Make sure to use the correct number of input channels to your network!!!\n",
    "\n",
    "Click [here](#seventh-todo) to go to the previous task.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd04855a",
   "metadata": {},
   "outputs": [],
   "source": [
    "prediction_type = 'affs'\n",
    "params = get_hyperparams(prediction_type)\n",
    "\n",
    "torch.manual_seed(42)\n",
    "\n",
    "d_factors = [[2,2],[2,2],[2,2]]\n",
    "in_channels= ... # set correct number of input channels\n",
    "num_fmaps=32\n",
    "fmap_inc_factors=3\n",
    "padding='same'\n",
    "final_kernel_size=1\n",
    "out_channels = params['out_channels']\n",
    "\n",
    "unet = UNet(\n",
    "        in_channels=in_channels,\n",
    "        num_fmaps=num_fmaps,\n",
    "        fmap_inc_factors=fmap_inc_factors,\n",
    "        downsample_factors=d_factors,\n",
    "        padding=padding)\n",
    "\n",
    "final_conv = torch.nn.Conv2d(\n",
    "    in_channels=num_fmaps,\n",
    "    out_channels=out_channels,\n",
    "    kernel_size=final_kernel_size)\n",
    "\n",
    "net = torch.nn.Sequential(unet, final_conv)\n",
    "\n",
    "net = net.to(device)\n",
    "\n",
    "summary(net, (in_channels, 128, 128))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "680a3216",
   "metadata": {},
   "outputs": [],
   "source": [
    "training_steps = 3000\n",
    "logdir = os.path.join(\"logs\", datetime.datetime.now().strftime(\"%Y%m%d-%H%M%S\"))\n",
    "writer = SummaryWriter(logdir)\n",
    "\n",
    "net = net.to(device)\n",
    "loss_fn = params['loss_function'].to(device)\n",
    "activation = params['activation']\n",
    "dtype = params['dtype']\n",
    "\n",
    "crop_size = 128\n",
    "\n",
    "### create datasets\n",
    "train_dataset = TissueNetDataset(root_dir='woodshole', split='train', crop_size=crop_size, prediction_type=prediction_type)\n",
    "test_dataset = TissueNetDataset(root_dir='woodshole', split='test', prediction_type=prediction_type)\n",
    "val_dataset = TissueNetDataset(root_dir='woodshole', split='test', val_split=True, prediction_type=prediction_type)\n",
    "\n",
    "batch_size = 4\n",
    "\n",
    "# make dataloaders\n",
    "train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, pin_memory=True)\n",
    "test_loader = DataLoader(test_dataset, batch_size=1)\n",
    "val_loader = DataLoader(val_dataset, batch_size=1)\n",
    "\n",
    "# set optimizer\n",
    "learning_rate = 1e-4\n",
    "optimizer = torch.optim.Adam(net.parameters(), lr=learning_rate)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf010db7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# run training loop\n",
    "train(train_loader, val_loader, net, loss_fn, activation, optimizer, dtype, prediction_type)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e630227",
   "metadata": {
    "tags": []
   },
   "source": [
    "<h3>Predict and visualize</h3>\n",
    "\n",
    "- Once the model is trained, we loop over the test dataset and visualize the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02cb28ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "avg = 0.0\n",
    "\n",
    "net.eval()\n",
    "\n",
    "for idx, (image, mask) in enumerate(test_loader):\n",
    "    image = image.to(device)\n",
    "    logits = net(image)\n",
    "    pred = activation(logits)\n",
    "        \n",
    "    image = np.squeeze(image.cpu())\n",
    "    mask = np.squeeze(mask.cpu().numpy())\n",
    "        \n",
    "    pred = np.squeeze(pred.cpu().detach().numpy())\n",
    "        \n",
    "    thresh = np.mean(pred)\n",
    "            \n",
    "    boundary_mask = get_boundary_mask(pred, prediction_type, thresh)\n",
    "    \n",
    "    if prediction_type == \"affs\":\n",
    "        boundary_distances = distance_transform_edt(boundary_mask)\n",
    "\n",
    "        pred_labels = watershed_from_boundary_distance(\n",
    "            boundary_distances,\n",
    "            boundary_mask\n",
    "        )\n",
    "        \n",
    "    else:\n",
    "        pred_labels = relabel_cc(boundary_mask)\n",
    "    \n",
    "    gt_labels = imread(test_loader.dataset.mask_files[idx])\n",
    "    \n",
    "    ap, precision, recall, tp, fp, fn = evaluate(gt_labels, pred_labels)\n",
    "    \n",
    "    avg += ap\n",
    "    \n",
    "    print(\n",
    "        f'Average precision: {ap} \\n',\n",
    "        f'Precision: {precision} \\n',\n",
    "        f'Recall: {recall} \\n',\n",
    "        f'True positives: {tp} \\n',\n",
    "        f'False positives: {fp} \\n',\n",
    "        f'False negatives: {fn} \\n'\n",
    "    )\n",
    "    \n",
    "    image = np.vstack((image, np.zeros_like(image)[:1]))\n",
    "    image = image.transpose(1,2,0)\n",
    "        \n",
    "    fig, axes = plt.subplots(1,3,figsize=(20, 20),sharex=True,sharey=True,squeeze=False)\n",
    "    \n",
    "    axes[0][0].imshow(image)\n",
    "    axes[0][0].title.set_text('Raw')\n",
    "    \n",
    "    axes[0][1].imshow(create_lut(gt_labels))\n",
    "    axes[0][1].title.set_text('GT Labels')\n",
    "    \n",
    "    axes[0][2].imshow(create_lut(pred_labels))\n",
    "    axes[0][2].title.set_text('Predicted Labels')\n",
    "    \n",
    "    plt.show()\n",
    "        \n",
    "avg /= (idx+1)\n",
    "    \n",
    "print(\"average precision on test set: {}\".format(avg))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97ea658f-230c-4554-94ea-a719368cc1da",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "<h3>Improve accuracy / segment cyto (time permitting</h3>\n",
    "    \n",
    "- It is likely that even with the extra things we added, we still aren't achieving the level of accuracy we would like to see. Production level models are usually trained for many more iterations and use lots of tricks to maximize the accuracy (also there are some false labels in the ground truth which will punish our network predictions even if correct). It is more important for now to conceptually understand the basics of instance segmentation and different approaches to increasing model robustness.\n",
    "    \n",
    "    \n",
    "- If you have time (now or in the future), try to improve the accuracy of your model. How accurate can you get it? Can you also get an accurate model on the cytoplasm masks? \n",
    "    \n",
    "    \n",
    "- If you have time (now or in the future), try the following bonus exercises - which show more advanced approaches to getting good instance segmentation results"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b1d5765",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "<hr style=\"height:2px;\"><div class=\"alert alert-block alert-success\"><h2>Checkpoint 3</h2>\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f63a89b",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Further learning\n",
    "\n",
    "* Instance segmentation can be challenging and this exercise just scratches the surface of what is possible.\n",
    "\n",
    "* There are bonus exercises (without todos) in the `bonus_exercises.ipynb` notebook. There are examples for CellPose and Local Shape Descriptors\n",
    "\n",
    "* This notebook assumes images that fit into memory but often times this is not the case (especially in biology). \n",
    "    1. To see an example for predicting over an image in chunks and stitching the results together, see this [notebook](https://github.com/dlmbl/instance_segmentation/blob/2021-solutions/3_tile_and_stitch.ipynb)\n",
    "    2. For a more advanced library that makes it easier to do machine learning on massive datasets, see gunpowder (navigate to the tutorials, or browse the API): https://funkelab.github.io/gunpowder\n",
    "    \n",
    "    \n",
    "* We did not cover more complex loss functions. Here are some nice explanations / implementations of other loss functions that are useful for instance segmentation: https://www.kaggle.com/code/bigironsphere/loss-function-library-keras-pytorch/notebook\n",
    "\n",
    "\n",
    "* A more complex (but powerful) approach is called metric learning. This can be seen in last years [exercise](https://github.com/dlmbl/instance_segmentation/blob/2021-solutions/2_instance_segmentation.ipynb)\n",
    "\n",
    "\n",
    "* We did not cover stardist in this tutorial, and barely scratched the surface on cellpose and lsds. For more tutorials on:\n",
    "    1. Stardist: https://github.com/maweigert/tutorials/tree/main/stardist\n",
    "    2. CellPose: https://github.com/MouseLand/cellpose#run-cellpose-10-without-local-python-installation\n",
    "    3. LSDs: https://github.com/funkelab/lsd#notebooks\n",
    "    \n",
    "## Good luck on your instance segmentation endeavors!!"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:06_instance_segmentation]",
   "language": "python",
   "name": "conda-env-06_instance_segmentation-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}