solutions.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "95227378-32db-418d-aa85-bf29eb8ad145",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Instance Segmentation\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 each TODO can link to the previous and next TODO to make it easier to move around freely. **Note:** This only works for uncollapsed cells. If you collapse a cell containing TODO 1 and then try to return to it using the TODO 2 link it won't work.\n",
    "\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\"><h1>Start here (AKA checkpoint 0)</h1>\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8812ec06-d766-4809-8d5f-aade999bbbc4",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Exercise 0.0: 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. We will import some more specific packages later."
   ]
  },
  {
   "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",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "from torchsummary import summary\n",
    "from tqdm.auto import tqdm\n",
    "from unet import *\n",
    "\n",
    "torch.backends.cudnn.benchmark = True"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c7bedd4-985d-41a9-b4c6-7c5669b716f0",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Exercise 1.0: 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": [
    "<div class=\"alert alert-block alert-info\"><h4>Exercise 1.1: Load and visualize data</h4>\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",
    "</div>\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": "e898df34-8889-44bd-9b4e-4917d497b9d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# utility function to view labels as rgb lut with matplotlib\n",
    "# eg plt.imshow(create_lut(labels))\n",
    "\n",
    "def create_lut(labels):\n",
    "\n",
    "    max_label = np.max(labels)\n",
    "\n",
    "    lut = np.random.randint(\n",
    "            low=0,\n",
    "            high=255,\n",
    "            size=(int(max_label + 1), 3),\n",
    "            dtype=np.uint8)\n",
    "\n",
    "    lut = np.append(\n",
    "            lut,\n",
    "            np.zeros(\n",
    "                (int(max_label + 1), 1),\n",
    "                dtype=np.uint8) + 255,\n",
    "            axis=1)\n",
    "\n",
    "    lut[0] = 0\n",
    "    colored_labels = lut[labels]\n",
    "\n",
    "    return colored_labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb7fbe07-8be2-426f-b256-cc8a20e48b20",
   "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",
    "# 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": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "<hr style=\"height:2px;\"><div class=\"alert alert-block alert-info\"><h4>Exercise 1.2: Create dataset</h4>\n",
    "    \n",
    "- Lets start by creating 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 like above. We will just use the first channel of the raw data (nuclei) for now (`raw[0]`)\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f61492b1",
   "metadata": {
    "tags": []
   },
   "source": [
    "<a id='first_todo'></a>\n",
    "\n",
    "##### **TODO (1)**\n",
    "\n",
    "* Create a basic dataset called `TissueNetDataset`\n",
    "* The `__init__` method should load our **mask file names** and **raw file names** into sorted lists\n",
    "* We will include optional arguments to make train/test splits and which mask files to choose\n",
    "* The `__len__` method should return the number of files in a given split folder (eg 50 for train)\n",
    "* The `__getitem__` method should get a single raw and mask file given the index and load into a numpy array.\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",
    "* Click [here](#second_todo) if you need to go to the next **TODO**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5c82461",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TissueNetDataset(Dataset):\n",
    "    def __init__(self,\n",
    "                 root_dir,\n",
    "                 split='train',\n",
    "                 mask='nuclei'\n",
    "                ):\n",
    "        \n",
    "        # use the root dir, split and mask logic to point to the correct files\n",
    "        self.mask_files = ... # load mask files into sorted list\n",
    "        self.raw_files = ... # load raw files into sorted list\n",
    "                \n",
    "    def __len__(self):\n",
    "        return ... # number of files\n",
    "        \n",
    "    def __getitem__(self, idx):\n",
    "        \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",
    "        return raw, mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "022e87f3-b13f-4934-a1b6-f2cf18005eb6",
   "metadata": {},
   "outputs": [],
   "source": [
    "#### Solution #####\n",
    "\n",
    "class TissueNetDataset(Dataset):\n",
    "    def __init__(self,\n",
    "                 root_dir,\n",
    "                 split='train',\n",
    "                 mask='nuclei'\n",
    "                ):\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",
    "    def __len__(self):\n",
    "        return len(self.raw_files)\n",
    "        \n",
    "    def __getitem__(self, idx):\n",
    "        # get file names\n",
    "        raw_file = self.raw_files[idx]\n",
    "        mask_file = self.mask_files[idx]\n",
    "        \n",
    "        # read in files\n",
    "        raw = imread(raw_file)\n",
    "        mask = imread(mask_file)\n",
    "\n",
    "        # for now just do single channel training\n",
    "        raw = raw[0]\n",
    "\n",
    "        return raw, mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d53c7c52-8604-4e72-9e48-da4e085e8b75",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = TissueNetDataset(root_dir='woodshole', split='train')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89ee65a7-75ee-4cbe-9443-2eab4d5cb8a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# visualize random batch\n",
    "\n",
    "raw, mask = train_dataset[random.randrange(len(train_dataset))]\n",
    "    \n",
    "plt.imshow(raw, cmap='gray')\n",
    "plt.imshow(create_lut(mask), alpha=0.5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d9ccd55-0b46-4e15-b348-82109155eade",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "<hr style=\"height:2px;\"><div class=\"alert alert-block alert-info\"><h4>Exercise 1.3: Add augmentations</h4>\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",
    "- Friendly reminder that raw[0] corresponds to nuclei, while raw[1] corresponds to cytoplasm :)\n",
    "- Below is a simple example. \n",
    "\n",
    "</div>\n",
    "\n",
    "![example_augmentation](static/example_augmentation.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2e7742a-f1c1-4804-989d-17cbc29761b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "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",
    "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": "f7d12135",
   "metadata": {},
   "source": [
    "<a id='second_todo'></a>\n",
    "\n",
    "##### **TODO (2)**\n",
    "\n",
    "* Add a parameter to your `TissueNetDataset` to define a `crop_size` of `64 pixels`\n",
    "* Add a function `augment_data` which builds our augmentations as shown above\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",
    "* Click [here](#first_todo) if you need to go back to the previous **TODO**\n",
    "* Click [here](#third_todo) if you need to go to the next **TODO**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ff3023c",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TissueNetDataset(Dataset):\n",
    "    def __init__(self,\n",
    "                 root_dir,\n",
    "                 split='train',\n",
    "                 mask='nuclei',\n",
    "                ):\n",
    "        \n",
    "        # make sure to add your split method since we will use it in `__getitem__`\n",
    "        # make sure to add your crop size\n",
    "        \n",
    "        # use the root dir, split and mask logic to point to the correct files\n",
    "        self.mask_files = ... # load mask files into sorted list\n",
    "        self.raw_files = ... # load raw files into sorted list\n",
    "        \n",
    "        \n",
    "    def __len__(self):\n",
    "        return len(self.raw_files)\n",
    "    \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",
    "        return raw, mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "249177e4-dd5f-485f-9014-c5050383c9e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "##### Solution #####\n",
    "\n",
    "class TissueNetDataset(Dataset):\n",
    "    def __init__(self,\n",
    "                 root_dir,\n",
    "                 split='train',\n",
    "                 mask='nuclei',\n",
    "                 crop_size=None\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",
    "        \n",
    "    def __len__(self):\n",
    "        return len(self.raw_files)\n",
    "    \n",
    "    def augment_data(self, raw, mask):\n",
    "        \n",
    "        transform = A.Compose([\n",
    "              A.RandomCrop(width=self.crop_size, height=self.crop_size),\n",
    "              A.HorizontalFlip(p=0.5),\n",
    "              A.VerticalFlip(p=0.5)\n",
    "            ])\n",
    "\n",
    "        transformed = transform(image=raw, mask=mask)\n",
    "\n",
    "        raw, mask = transformed['image'], transformed['mask']\n",
    "        \n",
    "        return raw, mask\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        # get file names\n",
    "        raw_file = self.raw_files[idx]\n",
    "        mask_file = self.mask_files[idx]\n",
    "        \n",
    "        # read in files\n",
    "        raw = imread(raw_file)\n",
    "        mask = imread(mask_file)\n",
    "\n",
    "        # for now just do single channel training\n",
    "        raw = raw[0]\n",
    "\n",
    "        # if training, run augmentations\n",
    "        if self.split == 'train':\n",
    "            raw, mask = self.augment_data(raw, mask)\n",
    "            \n",
    "        return raw, mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3004c9fd-0fb9-44f0-aa68-b53d196b6abd",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = TissueNetDataset(root_dir='woodshole', split='train', crop_size=64)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1016b12-97c8-4213-8028-539ddd93b4a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# visualize random batch\n",
    "\n",
    "raw, mask = train_dataset[random.randrange(len(train_dataset))]\n",
    "    \n",
    "plt.imshow(raw, cmap='gray')\n",
    "plt.imshow(create_lut(mask), alpha=0.5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be85bff2-da9a-42db-b7d1-28232c1f0bc0",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "<hr style=\"height:2px;\"><div class=\"alert alert-block alert-info\"><h4>Exercise 1.4: Create fg/bg representation</h4>\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",
    "</div>\n",
    "\n",
    "![example_fgbg](static/example_fgbg.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94babb25-5cce-498a-921e-212341fa55a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# function to erode boundary pixels. We fill an array with zeros, iterate over our labels, erode them\n",
    "# and then write them back into our empty array\n",
    "from scipy.ndimage import binary_erosion\n",
    "\n",
    "def erode(labels, iterations, border_value):\n",
    "\n",
    "    # copy labels to memory, create border array\n",
    "    labels = np.copy(labels)\n",
    "\n",
    "    # create zeros array for foreground\n",
    "    foreground = np.zeros_like(labels, dtype=bool)\n",
    "\n",
    "    # loop through unique labels\n",
    "    for label in np.unique(labels):\n",
    "\n",
    "        # skip background\n",
    "        if label == 0:\n",
    "            continue\n",
    "\n",
    "        # mask to label\n",
    "        label_mask = labels == label\n",
    "\n",
    "        # erode labels\n",
    "        eroded_mask = binary_erosion(\n",
    "                label_mask,\n",
    "                iterations=iterations,\n",
    "                border_value=border_value)\n",
    "\n",
    "        # get foreground\n",
    "        foreground = np.logical_or(eroded_mask, foreground)\n",
    "\n",
    "    # and background...\n",
    "    background = np.logical_not(foreground)\n",
    "\n",
    "    # set eroded pixels to zero\n",
    "    labels[background] = 0\n",
    "\n",
    "    return labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9450189a-647f-43e4-bda1-62e203653186",
   "metadata": {},
   "outputs": [],
   "source": [
    "# visualize the representation (repeatedly run cell)\n",
    "\n",
    "raw, mask = train_dataset[random.randrange(len(train_dataset))]\n",
    "\n",
    "labels = erode(\n",
    "    mask,\n",
    "    iterations=1,\n",
    "    border_value=1)\n",
    "\n",
    "labels_two_class = (labels != 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, cmap='gray')\n",
    "axes[0][0].title.set_text('Raw')\n",
    "\n",
    "axes[0][0].imshow(create_lut(mask), alpha=0.5)\n",
    "axes[0][0].title.set_text('Segmentation')\n",
    "\n",
    "axes[0][1].imshow(labels_two_class)\n",
    "axes[0][1].title.set_text('Foreground / background')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7cf7261",
   "metadata": {},
   "source": [
    "<a id='third_todo'></a>\n",
    "\n",
    "##### **TODO (3)**\n",
    "\n",
    "* Add a boolean parameter to your `TissueNetDataset` to create a validation split `val_split`\n",
    "* We can then use this to split our test data into validation / testing\n",
    "* After augmenting your data, create a fg / bg representation\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",
    "* Click [here](#second_todo) if you need to go back to the previous **TODO**\n",
    "* Click [here](#fourth_todo) if you need to go to the next **TODO**"
   ]
  },
  {
   "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",
    "                ):\n",
    "        \n",
    "        # make sure to add your split method since we will use it in `__getitem__`\n",
    "        # make sure to add your crop size\n",
    "        \n",
    "        # use the root dir, split and mask logic to point to the correct files\n",
    "        self.mask_files = ... # load mask files into sorted list\n",
    "        self.raw_files = ... # load raw 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\n",
    "        mask = ... # add channel dimension\n",
    "            \n",
    "        return raw, mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "952e5c57-0571-47e3-ab7c-a20f37605562",
   "metadata": {},
   "outputs": [],
   "source": [
    "#### Solution ####\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",
    "                ):\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",
    "        \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",
    "    def __len__(self):\n",
    "        return len(self.raw_files)\n",
    "    \n",
    "    def augment_data(self, raw, mask):\n",
    "        \n",
    "        transform = A.Compose([\n",
    "              A.RandomCrop(width=self.crop_size, height=self.crop_size),\n",
    "              A.HorizontalFlip(p=0.5),\n",
    "              A.VerticalFlip(p=0.5)\n",
    "        ])\n",
    "\n",
    "        transformed = transform(image=raw, mask=mask)\n",
    "\n",
    "        raw, mask = transformed['image'], transformed['mask']\n",
    "        \n",
    "        return raw, mask\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",
    "        # for now just do single channel training\n",
    "        raw = raw[0]\n",
    "\n",
    "        if self.split == 'train':\n",
    "            raw, mask = self.augment_data(raw, mask)\n",
    "            \n",
    "        fg = erode(\n",
    "                    mask,\n",
    "                    iterations=1,\n",
    "                    border_value=1)\n",
    "        \n",
    "        mask = (fg != 0).astype(np.float32)\n",
    "\n",
    "        # add channel dim for network\n",
    "        raw = np.expand_dims(raw, axis=0)\n",
    "        mask = np.expand_dims(mask, axis=0)\n",
    "                \n",
    "        return raw, mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42a94aef-8e0b-4d46-b2e2-6e476b6bca67",
   "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": "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(np.squeeze(raw), cmap='gray')\n",
    "axes[0][0].title.set_text('Raw')\n",
    "\n",
    "axes[0][1].imshow(np.squeeze(mask))\n",
    "axes[0][1].title.set_text('Mask')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "594c20e0-4cc6-419f-a1d8-8290ba39aa9e",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "<hr style=\"height:2px;\"><div class=\"alert alert-block alert-info\"><h4>Exercise 1.5: Create shallow network, visualize receptive field</h4>\n",
    "    \n",
    "- Let's create a shallow two layer network and visualize the receptive field. We will see later how this receptive field changes as we add more layers 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 layer -- 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",
    "</div>\n",
    "\n",
    "![example_RF](static/example_RF.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d5b6c59-3b2e-4caa-92cc-48162a80908d",
   "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]]\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=(5, 5))\n",
    "colors = [\"yellow\", \"red\", \"green\"]\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": "55352c82-130c-4452-badc-5be3e7ec6a7f",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "<hr style=\"height:2px;\"><div class=\"alert alert-block alert-info\"><h4>Exercise 1.6: Set hyperparameters, create model</h4>\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. \n",
    "    \n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02950213",
   "metadata": {},
   "source": [
    "<a id='fourth_todo'></a>\n",
    "\n",
    "##### **TODO (4)**  \n",
    "    \n",
    "- Decide how many output channels to have, remember we are starting with a binary task\n",
    "- What loss function and final layer 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 layer Unet with the following parameters:\n",
    "    - downsample by a factor of 2 in each layer\n",
    "    - single input channel\n",
    "    - 32 input feature maps \n",
    "    - multiply by a factor of 2 between layers\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",
    "- Click [here](#third_todo) if you need to go back to the previous **TODO**\n",
    "- Click [here](#fifth_todo) if you need to go to the next **TODO**"
   ]
  },
  {
   "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",
    "d_factors = ... \n",
    "in_channels = ... \n",
    "num_fmaps = ... \n",
    "fmap_inc_factors = ... \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",
    "device = ... # create your device. it should be cuda if a gpu is available, else cpu\n",
    "\n",
    "net = ... # cast your net to your device\n",
    "\n",
    "summary(net, (in_channels, 384, 384))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7cd42c5-6be2-49ca-987c-b0c7ccc40e73",
   "metadata": {},
   "outputs": [],
   "source": [
    "#### Solution ####\n",
    "\n",
    "out_channels = 1\n",
    "activation = torch.nn.Sigmoid()\n",
    "loss_fn = torch.nn.BCELoss()\n",
    "dtype = torch.FloatTensor\n",
    "\n",
    "torch.manual_seed(42)\n",
    "\n",
    "d_factors = [[2,2],[2,2]]\n",
    "in_channels=1\n",
    "num_fmaps=32\n",
    "fmap_inc_factors=2\n",
    "padding='same'\n",
    "final_kernel_size=1\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",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "net = net.to(device)\n",
    "\n",
    "summary(net, (in_channels, 384, 384))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b62b36d-00f0-4812-bf59-a94a6b9eebed",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "<hr style=\"height:2px;\"><div class=\"alert alert-block alert-info\"><h4>Exercise 1.7: Run training loop</h4>\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",
    "</div>"
   ]
  },
  {
   "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\n",
    "    logits = model(feature)\n",
    "    \n",
    "    if prediction_type == \"three_class\":\n",
    "        label=torch.squeeze(label,1)\n",
    "        \n",
    "    # final activation\n",
    "    predicted = activation(logits)\n",
    "\n",
    "    # pass through loss\n",
    "    loss_value = loss_fn(input=predicted, target=label)\n",
    "    \n",
    "    # backward if training mode\n",
    "    if train_step:\n",
    "        loss_value.backward()\n",
    "        optimizer.step()\n",
    "\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",
    "    with tqdm(total=training_steps) 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",
    "\n",
    "                    print(np.mean(acc_loss))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c734f964",
   "metadata": {},
   "source": [
    "<a id='fifth_todo'></a>\n",
    "\n",
    "##### **TODO (5)**\n",
    "\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",
    "- Click [here](#fourth_todo) if you need to go back to the previous **TODO**\n",
    "- Click [here](#sixth_todo) if you need to go to the next **TODO**"
   ]
  },
  {
   "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\n",
    "\n",
    "# run training loop... (eg call train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d81f44fa-e9f8-4568-89d5-74389252c834",
   "metadata": {},
   "outputs": [],
   "source": [
    "#### Solution ####\n",
    "\n",
    "train_batch_size = 4\n",
    "test_batch_size = 1\n",
    "\n",
    "# make dataloaders\n",
    "train_loader = DataLoader(train_dataset, batch_size=train_batch_size, shuffle=True, pin_memory=True)\n",
    "test_loader = DataLoader(test_dataset, batch_size=test_batch_size)\n",
    "val_loader = DataLoader(val_dataset, batch_size=test_batch_size)\n",
    "\n",
    "training_steps = 1000\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 = 1e-4\n",
    "optimizer = torch.optim.Adam(net.parameters(), lr=learning_rate)\n",
    "\n",
    "# run training loop\n",
    "train(train_loader, val_loader, net, loss_fn, activation, optimizer, dtype)"
   ]
  },
  {
   "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 forward the tensorflow port (usually 6006)\n",
    "# you can also do this by passing the host (relevant machine ip address), eg:\n",
    "\n",
    "# !tensorboard --logdir=logs --host hostname"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bba0107c-5a21-421f-8c33-fda6da798048",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "<hr style=\"height:2px;\"><div class=\"alert alert-block alert-info\"><h4>Exercise 1.8: Visualize results</h4>\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\n",
    "  \n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "daee8d9f-99b6-42fd-87a4-4e3c54c0ee0b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# convenience functions to threshold mask and relabel connected components as unique objects\n",
    "from skimage.filters import threshold_otsu\n",
    "from skimage.measure import label as relabel_cc\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\"><h1>Checkpoint 1</h1>\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de5f9fbe-33b7-4993-8060-91f2c44e7445",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Exercise 2.0: 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": "340d8bc9-825f-4c44-91a7-c08a0206bca9",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "<div class=\"alert alert-block alert-info\"><h4>Exercise 2.1: Add extra representations</h4>\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",
    "</div>\n",
    "\n",
    "![diff_costs](static/diff_costs.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e63984d-53a6-4e21-a2c6-bea4216d57f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# convenience functions to create our three extra representations (three class, sdt, affs)\n",
    "# run this cell to initialize these functions for our dataset\n",
    "from scipy.ndimage import distance_transform_edt\n",
    "\n",
    "# utility function to compute a three class representation\n",
    "def erode(labels, iterations, border_value):\n",
    "\n",
    "    # copy labels to memory, create border array\n",
    "    labels = np.copy(labels)\n",
    "    border = np.array(labels)\n",
    "\n",
    "    # create zeros array for foreground\n",
    "    foreground = np.zeros_like(labels, dtype=bool)\n",
    "\n",
    "    # loop through unique labels\n",
    "    for label in np.unique(labels):\n",
    "\n",
    "        # skip background\n",
    "        if label == 0:\n",
    "            continue\n",
    "\n",
    "        # mask to label\n",
    "        label_mask = labels == label\n",
    "\n",
    "        # erode labels\n",
    "        eroded_mask = binary_erosion(\n",
    "                label_mask,\n",
    "                iterations=iterations,\n",
    "                border_value=border_value)\n",
    "\n",
    "        # get foreground\n",
    "        foreground = np.logical_or(eroded_mask, foreground)\n",
    "\n",
    "    # and background...\n",
    "    background = np.logical_not(foreground)\n",
    "\n",
    "    # set eroded pixels to zero\n",
    "    labels[background] = 0\n",
    "\n",
    "    # get eroded pixels\n",
    "    border = labels - border\n",
    "\n",
    "    return labels, border\n",
    "\n",
    "\n",
    "# utility function to compute a signed distance transform\n",
    "def compute_sdt(labels, constant=0.5, scale=5):\n",
    "\n",
    "    inner = distance_transform_edt(binary_erosion(labels))\n",
    "    outer = distance_transform_edt(np.logical_not(labels))\n",
    "\n",
    "    distance = (inner - outer) + constant\n",
    "\n",
    "    distance = np.tanh(distance / scale)\n",
    "\n",
    "    return distance\n",
    "\n",
    "\n",
    "# utility function to compute edge affinities\n",
    "def compute_affinities(seg, nhood):\n",
    "\n",
    "    nhood = np.array(nhood)\n",
    "\n",
    "    shape = seg.shape\n",
    "    nEdge = nhood.shape[0]\n",
    "    dims = nhood.shape[1]\n",
    "    aff = np.zeros((nEdge,) + shape, dtype=np.int32)\n",
    "\n",
    "    for e in range(nEdge):\n",
    "        aff[e, \\\n",
    "          max(0,-nhood[e,0]):min(shape[0],shape[0]-nhood[e,0]), \\\n",
    "          max(0,-nhood[e,1]):min(shape[1],shape[1]-nhood[e,1])] = \\\n",
    "                      (seg[max(0,-nhood[e,0]):min(shape[0],shape[0]-nhood[e,0]), \\\n",
    "                          max(0,-nhood[e,1]):min(shape[1],shape[1]-nhood[e,1])] == \\\n",
    "                        seg[max(0,nhood[e,0]):min(shape[0],shape[0]+nhood[e,0]), \\\n",
    "                          max(0,nhood[e,1]):min(shape[1],shape[1]+nhood[e,1])] ) \\\n",
    "                      * ( seg[max(0,-nhood[e,0]):min(shape[0],shape[0]-nhood[e,0]), \\\n",
    "                          max(0,-nhood[e,1]):min(shape[1],shape[1]-nhood[e,1])] > 0 ) \\\n",
    "                      * ( seg[max(0,nhood[e,0]):min(shape[0],shape[0]+nhood[e,0]), \\\n",
    "                          max(0,nhood[e,1]):min(shape[1],shape[1]+nhood[e,1])] > 0 )\n",
    "                          \n",
    "\n",
    "    return aff"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70f98e43-335a-40c3-8363-6ad6a1d9ca6e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 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",
    "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",
    "labels, border = erode(\n",
    "    aug_mask,\n",
    "    iterations=1,\n",
    "    border_value=1)\n",
    "\n",
    "labels_two_class = (labels != 0)\n",
    "border[border!=0] = 2\n",
    "\n",
    "labels_three_class = (labels_two_class + border)\n",
    "sdt = compute_sdt(labels)\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": {},
   "source": [
    "<a id='sixth_todo'></a>\n",
    "\n",
    "##### **TODO (6)**\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",
    "- Your `__getitem__` should call `create_target` after augmenting your data\n",
    "- Click [here](#fifth_todo) if you need to go back to the previous **TODO**\n",
    "- Click [here](#seventh_todo) if you need to go to the next **TODO**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "045e96e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TissueNetDataset(Dataset):\n",
    "    def __init__(self,\n",
    "                 root_dir,\n",
    "                 split='train',\n",
    "                 mask='nuclei',\n",
    "                ):\n",
    "        \n",
    "        # make sure to add your split method since we will use it in `__getitem__`\n",
    "        # make sure to add your crop size\n",
    "        \n",
    "        # use the root dir, split and mask logic to point to the correct files\n",
    "        self.mask_files = ... # load mask files into sorted list\n",
    "        self.raw_files = ... # load raw files into sorted list\n",
    "        \n",
    "        # Add another parameter `val_split`. \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",
    "        # Add a parameter `prediction_type` to define which representation to choose.\n",
    "        # Default should be `two_class`\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",
    "    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\n",
    "        elif prediction_type == 'affs':\n",
    "            mask = ... # get affs\n",
    "            \n",
    "        mask = ... # cast mask to float32\n",
    "        \n",
    "        return mask  \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",
    "        raw, mask = ... # augment your data **if** split mode is train\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": "777bead0-3140-4887-8080-228c7d01d6aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "#### Solution ####\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",
    "                ):\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",
    "        \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",
    "    def __len__(self):\n",
    "        return len(self.raw_files)\n",
    "    \n",
    "    def augment_data(self, raw, mask):\n",
    "        \n",
    "        transform = A.Compose([\n",
    "              A.RandomCrop(width=self.crop_size, height=self.crop_size),\n",
    "              A.HorizontalFlip(p=0.5),\n",
    "              A.VerticalFlip(p=0.5)\n",
    "            ])\n",
    "\n",
    "        transformed = transform(image=raw, mask=mask)\n",
    "\n",
    "        raw, mask = transformed['image'], transformed['mask']\n",
    "        \n",
    "        return raw, mask\n",
    "    \n",
    "    def create_target(self, mask, prediction_type):\n",
    "        \n",
    "        mask, border = erode(\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",
    "    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",
    "        # for now just do single channel training\n",
    "        raw = raw[0]\n",
    "\n",
    "        if self.split == 'train':\n",
    "            raw, mask = self.augment_data(raw, mask)\n",
    "            \n",
    "        mask = self.create_target(mask, self.prediction_type)\n",
    "        \n",
    "        # add channel dim for network\n",
    "        raw = np.expand_dims(raw, axis=0)\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": "759137ac-7f13-4e48-bd75-954ec9f64f2d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# try each prediction type\n",
    "prediction_type = 'affs'\n",
    "\n",
    "train_dataset = TissueNetDataset(\n",
    "    root_dir='woodshole',\n",
    "    split='train',\n",
    "    crop_size=64,\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": "02aa001e-797d-4d6f-9735-b81ad1af2ddc",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "<hr style=\"height:2px;\"><div class=\"alert alert-block alert-info\"><h4>Exercise 2.2: Add more augmentations</h4>\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",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f28639b9",
   "metadata": {
    "tags": []
   },
   "source": [
    "<a id='seventh_todo'></a>\n",
    "\n",
    "##### **TODO (7)**\n",
    "\n",
    "- Update your augmentation function to use the following augmentations in order (use a 30% probability for any transform that requires it):\n",
    "    - RandomCrop() (by `crop_size`)\n",
    "    - PadIfNeeded() (see below for instructions)\n",
    "    - HorizontalFlip()\n",
    "    - VerticalFlip()\n",
    "    - RandomRotate90()\n",
    "    - Transpose()\n",
    "    - RandomBrightnessContrast\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",
    "    - 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",
    "    - 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",
    "    - Use a probability of 1 (the padding should always occur if needed)\n",
    "    - Use a `border_mode` of 0 to pad it with zeros (reflect is the default which we don't want)\n",
    "- Click [here](#sixth_todo) if you need to go back to the previous **TODO**\n",
    "- Click [here](#eighth_todo) if you need to go to the next **TODO**"
   ]
  },
  {
   "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",
    "                ):\n",
    "        \n",
    "        # make sure to add your split method since we will use it in `__getitem__`\n",
    "        # make sure to add your crop size\n",
    "        \n",
    "        # use the root dir, split and mask logic to point to the correct files\n",
    "        self.mask_files = ... # load mask files into sorted list\n",
    "        self.raw_files = ... # load raw files into sorted list\n",
    "        \n",
    "        # Add another parameter `val_split`. \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",
    "        # 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",
    "        \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",
    "        return padding\n",
    "    \n",
    "    \n",
    "    def augment_data(self, raw, mask):\n",
    "        \n",
    "        transform = ... # create your new augmentations\n",
    "        \n",
    "        transformed = ... # call your augmentations\n",
    "        \n",
    "        raw, mask = ... # get your resulting arrays\n",
    "        \n",
    "        return raw, mask\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\n",
    "        elif prediction_type == 'affs':\n",
    "            mask = ... # get affs\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 = ... # load raw to numpy array\n",
    "        mask = ... # load mask to numpy array\n",
    "        \n",
    "        raw = ... # get nuclei channel\n",
    "        \n",
    "        raw, mask = ... # get your padding **and** augment your data **if** split mode is train\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": "6bc7f252-29f4-4352-9f6f-fbb316fdbb46",
   "metadata": {},
   "outputs": [],
   "source": [
    "#### Solution ####\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",
    "    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 augment_data(self, raw, mask, padding):\n",
    "        \n",
    "        transform = A.Compose([\n",
    "              A.RandomCrop(\n",
    "                  width=self.crop_size,\n",
    "                  height=self.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\n",
    "    \n",
    "\n",
    "    def create_target(self, mask, prediction_type):\n",
    "        \n",
    "        mask, border = erode(\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",
    "        # for now just do single channel training\n",
    "        raw = raw[0]\n",
    "\n",
    "        if self.split == 'train':\n",
    "            padding = self.get_padding(self.crop_size, self.padding_size)\n",
    "            raw, mask = self.augment_data(raw, mask, padding)\n",
    "            \n",
    "        mask = self.create_target(mask, self.prediction_type)\n",
    "        \n",
    "        # add channel dim for network\n",
    "        raw = np.expand_dims(raw, axis=0)\n",
    "        \n",
    "        if self.prediction_type != 'affs':\n",
    "            mask = np.expand_dims(mask, axis=0)\n",
    "            \n",
    "        return raw, mask"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f38a4a0a-2d9e-4fbb-86bd-a592fb8bb994",
   "metadata": {},
   "source": [
    "* Try each prediction type \n",
    "* What happens if you increase crop size to 65? - how does the padding come into play?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aeced379-98f3-4d32-b4ba-c59cf579a748",
   "metadata": {},
   "outputs": [],
   "source": [
    "prediction_type = 'affs'\n",
    "crop_size = 64\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": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "<hr style=\"height:2px;\"><div class=\"alert alert-block alert-info\"><h4>Exercise 2.3: Add extra hyperparameters based on prediction type</h4>\n",
    "    \n",
    "- Now that we have more representations, we need to be sure that our hyperparams are consistent with whichever representation we choose. \n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "357d9786",
   "metadata": {
    "tags": []
   },
   "source": [
    "<a id='eighth_todo'></a>\n",
    "\n",
    "##### **TODO (8)**\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: `cross entropy` and `mean squared error`\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",
    "- Click [here](#seventh_todo) if you need to go back to the previous **TODO**\n",
    "- Click [here](#ninth_todo) if you need to go to the next **TODO**"
   ]
  },
  {
   "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 RuntimeError(\"invalid prediction type\")\n",
    "        \n",
    "    params = ... # get dict\n",
    "    \n",
    "    return params"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e27752c5-1d43-4560-830a-8681e90ec922",
   "metadata": {},
   "outputs": [],
   "source": [
    "#### Solution ####\n",
    "\n",
    "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",
    "        out_channels = 3\n",
    "        activation = torch.nn.Softmax(dim=1)\n",
    "        loss_fn = torch.nn.CrossEntropyLoss()\n",
    "        dtype = torch.LongTensor\n",
    "\n",
    "    elif prediction_type == \"sdt\":\n",
    "        out_channels = 1\n",
    "        activation = torch.nn.Tanh()\n",
    "        loss_fn = torch.nn.MSELoss()\n",
    "        dtype = torch.FloatTensor\n",
    "\n",
    "    elif prediction_type == \"affs\":\n",
    "        out_channels = 2\n",
    "        activation = torch.nn.Sigmoid()\n",
    "        loss_fn = torch.nn.MSELoss()\n",
    "        dtype = torch.FloatTensor\n",
    "\n",
    "    else:\n",
    "        raise RuntimeError(\"invalid prediction type\")\n",
    "        \n",
    "    params = {\n",
    "        'out_channels': out_channels,\n",
    "        'activation': activation,\n",
    "        'loss_function': loss_fn,\n",
    "        'dtype': dtype\n",
    "    }\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": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "<hr style=\"height:2px;\"><div class=\"alert alert-block alert-info\"><h4>Exercise 2.4: Increase batch crop size</h4>\n",
    "    \n",
    "- Before we were using a smaller batch 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",
    "\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": [
    "<hr style=\"height:2px;\"><div class=\"alert alert-block alert-info\"><h4>Exercise 2.5: Increase features in network</h4>\n",
    "    \n",
    "- Before we were training with a pretty small network, eg two downsampling layers. Let's see how the receptive fields change as we increase our network to 3 layers\n",
    "    \n",
    "    \n",
    "- Try changing the downsampling factors as before. How does it change the receptive field?\n",
    "\n",
    "</div>\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": {},
   "source": [
    "<a id='ninth_todo'></a>\n",
    "\n",
    "##### **TODO (9)**\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",
    "- Click [here](#eighth_todo) if you need to go back to the previous **TODO**\n",
    "- Click [here](#tenth_todo) if you need to go to the next **TODO**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac78b005-f557-48c2-a37a-f12f2847b0d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "#### Example Solution\n",
    "\n",
    "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=1\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",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "net = net.to(device)\n",
    "\n",
    "summary(net, (in_channels, 384, 384))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "017c5fef-7fec-49c5-878a-00ed2f81dc90",
   "metadata": {
    "tags": []
   },
   "source": [
    "<hr style=\"height:2px;\"><div class=\"alert alert-block alert-info\"><h4>Exercise 2.6: Train for longer</h4>\n",
    "    \n",
    "- Let's train our network for a decent amount of steps\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "869beb47",
   "metadata": {},
   "source": [
    "<a id='tenth_todo'></a>\n",
    "\n",
    "##### **TODO (10)**\n",
    "\n",
    "- Create your datasets and loaders as before. Increase your batch size (eg 64 -> 128)    \n",
    "- Use the same learning rate and optimizer as before   \n",
    "- Make sure you correctly set your loss_fn, activation, and tensor dtype (remember we now define these with our `get_hyperparams` function)\n",
    "- Train for longer (eg 1000 -> 3000 steps)\n",
    "- Click [here](#ninth_todo) if you need to go back to the previous **TODO**\n",
    "- Click [here](#eleventh_todo) if you need to go to the next **TODO**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7273eea-678f-497b-9b66-7d2ec96c9e26",
   "metadata": {},
   "outputs": [],
   "source": [
    "#### Solution ####\n",
    "\n",
    "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)\n",
    "\n",
    "# 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\"><h1>Checkpoint 2</h1>\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3ffca7a-7dce-476e-99f2-c351b02aaade",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Exercise 3.0: Post-processing / further improvements\n",
    "\n",
    "- Before we were just thresholding and relabeling connected components. Now we will see a more advanced post-processing strategy called watershed.\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": [
    "<div class=\"alert alert-block alert-info\"><h4>Exercise 3.1: Introduce watershed</h4>\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\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86e1af4e-44c1-48e3-8ba1-f7b4833880b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# utility function to compute watershed from a boundary mask created from the predictions\n",
    "from skimage.segmentation import watershed\n",
    "from scipy.ndimage import maximum_filter\n",
    "from scipy.ndimage import label as label_cc\n",
    "\n",
    "def watershed_from_boundary_distance(\n",
    "        boundary_distances,\n",
    "        boundary_mask,\n",
    "        id_offset=0,\n",
    "        min_seed_distance=10):\n",
    "\n",
    "    # get our seeds \n",
    "    max_filtered = maximum_filter(boundary_distances, min_seed_distance)\n",
    "    maxima = max_filtered==boundary_distances\n",
    "    seeds, n = label_cc(maxima)\n",
    "\n",
    "    if n == 0:\n",
    "        return np.zeros(boundary_distances.shape, dtype=np.uint64), id_offset\n",
    "\n",
    "    seeds[seeds!=0] += id_offset\n",
    "\n",
    "    # calculate our segmentation\n",
    "    segmentation = watershed(\n",
    "        boundary_distances.max() - boundary_distances,\n",
    "        seeds,\n",
    "        mask=boundary_mask)\n",
    "    \n",
    "    return segmentation\n",
    "\n",
    "def get_boundary_mask(pred, prediction_type, thresh=None):\n",
    "    \n",
    "    if prediction_type == 'two_class' or prediction_type == 'sdt':\n",
    "        # simple threshold\n",
    "        boundary_mask = pred > thresh\n",
    "\n",
    "    elif prediction_type == 'three_class':\n",
    "        # Return the indices of the maximum values along channel axis, then set mask to cell interior (1)\n",
    "        boundary_mask = np.argmax(pred, axis=0)\n",
    "        boundary_mask = boundary_mask == 1\n",
    "\n",
    "    elif prediction_type == 'affs':\n",
    "        # take mean of combined affs then threshold\n",
    "        boundary_mask = 0.5 * (pred[0] + pred[1]) > thresh\n",
    "    else:\n",
    "        raise Exception('Choose from one of the following prediction types: two_class, three_class, sdt, affs')\n",
    "        \n",
    "    return boundary_mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "692be5ec-8f3f-4ddb-9e28-02a554dfb023",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 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",
    "    boundary_mask = get_boundary_mask(pred, prediction_type, thresh=thresh)\n",
    "    boundary_distances = distance_transform_edt(boundary_mask)\n",
    "    \n",
    "    seg = watershed_from_boundary_distance(\n",
    "        boundary_distances,\n",
    "        boundary_mask\n",
    "    )\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": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "<hr style=\"height:2px;\"><div class=\"alert alert-block alert-info\"><h4>Exercise 3.2: Evaluate</h4>\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": "9fd80a34-f900-4c74-8ca8-1961b2d8f05e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Utility function to evaluate our segmentations\n",
    "\n",
    "from scipy.optimize import linear_sum_assignment\n",
    "from skimage.segmentation import relabel_sequential\n",
    "\n",
    "def evaluate(gt_labels, pred_labels, th=0.5):\n",
    "    pred_labels_rel, _, _ = relabel_sequential(pred_labels)\n",
    "    gt_labels_rel, _, _ = relabel_sequential(gt_labels)\n",
    "\n",
    "    overlay = np.array([pred_labels_rel.flatten(),\n",
    "                        gt_labels_rel.flatten()])\n",
    "\n",
    "    # get overlaying cells and the size of the overlap\n",
    "    overlay_labels, overlay_labels_counts = np.unique(\n",
    "        overlay, return_counts=True, axis=1)\n",
    "    overlay_labels = np.transpose(overlay_labels)\n",
    "\n",
    "    # get gt cell ids and the size of the corresponding cell\n",
    "    gt_labels_list, gt_counts = np.unique(gt_labels_rel, return_counts=True)\n",
    "    gt_labels_count_dict = {}\n",
    "    \n",
    "    for (l, c) in zip(gt_labels_list, gt_counts):\n",
    "        gt_labels_count_dict[l] = c\n",
    "\n",
    "    # get pred cell ids\n",
    "    pred_labels_list, pred_counts = np.unique(pred_labels_rel,\n",
    "                                              return_counts=True)\n",
    "\n",
    "    pred_labels_count_dict = {}\n",
    "    for (l, c) in zip(pred_labels_list, pred_counts):\n",
    "        pred_labels_count_dict[l] = c\n",
    "\n",
    "    num_pred_labels = int(np.max(pred_labels_rel))\n",
    "    num_gt_labels = int(np.max(gt_labels_rel))\n",
    "    num_matches = min(num_gt_labels, num_pred_labels)\n",
    "    \n",
    "    # create iou table\n",
    "    iouMat = np.zeros((num_gt_labels+1, num_pred_labels+1),\n",
    "                      dtype=np.float32)\n",
    "\n",
    "    for (u, v), c in zip(overlay_labels, overlay_labels_counts):\n",
    "        iou = c / (gt_labels_count_dict[v] + pred_labels_count_dict[u] - c)\n",
    "        iouMat[int(v), int(u)] = iou\n",
    "\n",
    "    # remove background\n",
    "    iouMat = iouMat[1:, 1:]\n",
    "\n",
    "    # use IoU threshold th\n",
    "    if num_matches > 0 and np.max(iouMat) > th:\n",
    "        costs = -(iouMat > th).astype(float) - iouMat / (2*num_matches)\n",
    "        gt_ind, pred_ind = linear_sum_assignment(costs)\n",
    "        assert num_matches == len(gt_ind) == len(pred_ind)\n",
    "        match_ok = iouMat[gt_ind, pred_ind] > th\n",
    "        tp = np.count_nonzero(match_ok)\n",
    "    else:\n",
    "        tp = 0\n",
    "    fp = num_pred_labels - tp\n",
    "    fn = num_gt_labels - tp\n",
    "    ap = tp / max(1, tp + fn + fp)\n",
    "    precision = tp / max(1, tp + fp)\n",
    "    recall = tp / max(1, tp + fn)\n",
    "\n",
    "    return ap, precision, recall, tp, fp, fn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0933b0ef-6731-4fc6-a746-21f45acde9d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 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",
    "    boundary_distances = distance_transform_edt(boundary_mask)\n",
    "    \n",
    "    pred_labels = watershed_from_boundary_distance(\n",
    "        boundary_distances,\n",
    "        boundary_mask\n",
    "    )\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": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "<hr style=\"height:2px;\"><div class=\"alert alert-block alert-info\"><h4>Exercise 3.3: Loop over batches</h4>\n",
    "    \n",
    "- Now we can loop over all test set images and get the the average model precision\n",
    "\n",
    "</div>"
   ]
  },
  {
   "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",
    "    boundary_distances = distance_transform_edt(boundary_mask)\n",
    "    \n",
    "    pred_labels = watershed_from_boundary_distance(\n",
    "        boundary_distances,\n",
    "        boundary_mask\n",
    "    )\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": [
    "<hr style=\"height:2px;\"><div class=\"alert alert-block alert-info\"><h4>Exercise 3.4: Use both channels of raw data</h4>\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",
    "    \n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c2c19d7",
   "metadata": {},
   "source": [
    "<a id='eleventh_todo'></a>\n",
    "\n",
    "##### **TODO (11)**\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",
    "- Click [here](#tenth_todo) if you need to go back to the previous **TODO**\n",
    "- Click [here](#twelfth_todo) if you need to go to the next **TODO**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f701c49",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TissueNetDataset(Dataset):\n",
    "    # Create your dataset :)\n",
    "    pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "350322ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "#### Solution ####\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",
    "    def __len__(self):\n",
    "        return len(self.raw_files)\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",
    "    def augment_data(self, raw, mask, padding):\n",
    "        \n",
    "        transform = A.Compose([\n",
    "              A.RandomCrop(\n",
    "                  width=self.crop_size,\n",
    "                  height=self.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\n",
    "    \n",
    "    def create_target(self, mask, prediction_type):\n",
    "        \n",
    "        mask, border = erode(\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",
    "    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 = raw.transpose([1,2,0])   \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 = self.augment_data(raw, mask, padding)\n",
    "            \n",
    "        raw = raw.transpose([2,0,1])\n",
    "        mask = mask.transpose([2,0,1])\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": {},
   "source": [
    "<a id='twelfth_todo'></a>\n",
    "\n",
    "##### **TODO (12)**   \n",
    "    \n",
    "- Create your network, hyperparameters and data loaders as last time (eg 3 layers, 3 mult factor, 128 crop size, 3k iterations, etc). Make sure to use the correct number of input channels to your network!!!\n",
    "- Click [here](#eleventh_todo) if you need to go back to the previous **TODO**\n",
    "- Click [here](#final_todo) if you need to go to the next **TODO**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76fae090",
   "metadata": {},
   "outputs": [],
   "source": [
    "#### Example Solution\n",
    "\n",
    "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=2\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",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "net = net.to(device)\n",
    "\n",
    "summary(net, (in_channels, 384, 384))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "680a3216",
   "metadata": {},
   "outputs": [],
   "source": [
    "#### Example Solution ####\n",
    "\n",
    "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)\n",
    "\n",
    "# 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": [
    "<a id='final_todo'></a>\n",
    "\n",
    "##### **TODO (13)**\n",
    "\n",
    "- Once your network is trained:\n",
    "    - loop over your test loader\n",
    "    - pass each image through your net and activation to get your predictions\n",
    "    - create a boundary mask and boundary distances from your prediction\n",
    "    - watershed\n",
    "    - evaluate\n",
    "    - visualize (remember you will need to visualize the raw data accordingly since we are now doing multi-channel)\n",
    "- Click [here](#twelfth_todo) if you need to go back to the previous **TODO**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02cb28ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "#### Solution ####\n",
    "\n",
    "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",
    "    boundary_distances = distance_transform_edt(boundary_mask)\n",
    "    \n",
    "    pred_labels = watershed_from_boundary_distance(\n",
    "        boundary_distances,\n",
    "        boundary_mask\n",
    "    )\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": [
    "<hr style=\"height:2px;\"><div class=\"alert alert-block alert-info\"><h4>Exercise 3.5 (time permitting): Improve accuracy / segment cyto</h4>\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\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b1d5765",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "<hr style=\"height:2px;\"><div class=\"alert alert-block alert-success\"><h1>Checkpoint 3</h1>\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b6ace09-2b95-4ad8-bde0-483486c16818",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "## Exercise 4.0: Bonus exercises\n",
    "\n",
    "- These exercises do not have todos, feel free to run them to get a sense of a few more tricks you can use for instance segmentation. They also might come in handy later when you are working on your projects"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b593025-1e9a-4cf1-96ef-c177a7c450a6",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "<div class=\"alert alert-block alert-info\"><h4>Exercise 4.1: Balance labels, use weighted loss</h4>\n",
    "    \n",
    "- Before we were treating every pixel as equally important, but this is not always the case. Let's look at the three class representation as an example (assumes you have already created the final TissueNetDataset correctly :))\n",
    "    \n",
    "    \n",
    "- We have significantly more background pixels (label 0) than border pixels (label 2). So why don't we weight them differently? We should have a higher weighting on the border pixels. \n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "324f4b45-e6b9-483f-8040-9dbceeba5768",
   "metadata": {},
   "outputs": [],
   "source": [
    "prediction_type = 'three_class'\n",
    "\n",
    "params = get_hyperparams(prediction_type)\n",
    "params"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b022d720-9487-4f40-b60c-f0832159567b",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = TissueNetDataset(\n",
    "    root_dir='woodshole',\n",
    "    split='train',\n",
    "    crop_size=128,\n",
    "    prediction_type=prediction_type)\n",
    "\n",
    "background = []\n",
    "inner = []\n",
    "border = []\n",
    "\n",
    "for raw, mask in train_dataset:\n",
    "\n",
    "    # get unique labels and how many pixels they have (this should just be 0,1,2 for three class)\n",
    "    labels, counts = np.unique(mask, return_counts=True)\n",
    "    \n",
    "    # test for background batch (some batches could be entirely background)\n",
    "    if len(counts) != 3:\n",
    "        continue\n",
    "    \n",
    "    background.append(counts[0])\n",
    "    inner.append(counts[1])\n",
    "    border.append(counts[2])\n",
    "    \n",
    "# get the averages across all training batches\n",
    "# to make sure we don't have a misrepresented weighting from a single batch\n",
    "all_counts = [np.mean(i) for i in (background, inner, border)]\n",
    "\n",
    "# get the probabilities as they sum up to 1 \n",
    "label_proportions = [(label, count/sum(counts)) for label,count in zip(labels, all_counts)]\n",
    "\n",
    "print(label_proportions)\n",
    "\n",
    "# our loss function can take a weight. this should be the inverse of the label proportions. \n",
    "# eg:\n",
    "    # if label 0 has a proportion of 0.55, it would have a weight of 1/0.55 = 1.8\n",
    "    # if label 2 has a proportion of 0.09 it would have a weight of 1/0.09 = 11\n",
    "\n",
    "params['loss_function'] = torch.nn.CrossEntropyLoss(\n",
    "    weight=torch.tensor(\n",
    "        [(1/p).astype(np.float32) for l, p in label_proportions]\n",
    "    ))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "588e1568-d59f-42c2-b2f4-b2e6662a08a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(42)\n",
    "\n",
    "d_factors = [[2,2],[2,2],[2,2]]\n",
    "\n",
    "in_channels=2\n",
    "num_fmaps=32\n",
    "fmap_inc_factors=3\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",
    "        activation='ReLU',\n",
    "        padding='same')\n",
    "\n",
    "final_conv = torch.nn.Conv2d(\n",
    "    in_channels=num_fmaps,\n",
    "    out_channels=out_channels,\n",
    "    kernel_size=1,\n",
    "    padding=0,\n",
    "    bias=True)\n",
    "\n",
    "net = torch.nn.Sequential(unet, final_conv)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a94dcce6-ec63-403f-ac41-bc676054163a",
   "metadata": {},
   "outputs": [],
   "source": [
    "training_steps = 2000\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",
    "# set optimizer\n",
    "learning_rate = 1e-4\n",
    "optimizer = torch.optim.Adam(net.parameters(), lr=learning_rate)\n",
    "\n",
    "### create datasets\n",
    "\n",
    "train_dataset = TissueNetDataset(root_dir='woodshole', split='train', crop_size=128, 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",
    "# run training loop\n",
    "train(train_loader, val_loader, net, loss_fn, activation, optimizer, dtype, prediction_type)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6fad097-c9a5-4932-b53c-670097d997e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "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",
    "    pred = np.argmax(pred, axis=0)\n",
    "        \n",
    "    fig, axes = plt.subplots(1,3,figsize=(20, 20),sharex=True,sharey=True,squeeze=False)\n",
    "    \n",
    "    image = np.vstack((image, np.zeros_like(image)[:1]))\n",
    "    image = image.transpose(1,2,0)\n",
    "    \n",
    "    axes[0][0].imshow(image)\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(pred)\n",
    "    axes[0][2].title.set_text('Predicted')\n",
    "    \n",
    "    if idx == 2:\n",
    "        break"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a5d4651b-2713-4c45-8aec-9925ea0b88b8",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "<hr style=\"height:2px;\"><div class=\"alert alert-block alert-info\"><h4>Exercise 4.2: Auxiliary learning</h4>\n",
    "    \n",
    "- Auxiliary learning is a powerful technique that can help to improve the results of our main objective by providing a helper task. Up until now, we have only shown our model representations of the data that are boundary specific. But the data is a lot richer than that - these objects have distinct shapes that could be leveraged in order to better learn the boundaries.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7fdcc839",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "##### **Cellpose**\n",
    "    \n",
    "- In [**Cellpose**](https://cellpose.readthedocs.io/en/latest/), cells are turned into flow representations. We create these flow representations by simulating diffusion from the center of the cell to get the spatial gradients for each pixel that point towards the center of the cell. During test time, we use the flows as a dynamical system and all pixels that converge to the same point are defined as the pixels in a given cell. The flows shown below are represented by an HSV colormap used in the optic flow literature.\n",
    "    \n",
    "    \n",
    "- We also predict the foregroud / background -- the two classes you predicted in exercise 1. In Cellpose we call this the cell probability. We threshold this to decide which pixels are in cells -- we only use these pixels to run the dynamical system.\n",
    "    \n",
    "    \n",
    "- The flow representation allows the learning of non-convex shapes, because pixels can flow around corners. It also prevents merging, as flows for two cells that are touching are opposite.\n",
    "\n",
    "![cellpose_flows](static/cellpose_flows.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e880baf",
   "metadata": {},
   "outputs": [],
   "source": [
    "from cellpose import models\n",
    "\n",
    "# create a cellpose model on the gpu\n",
    "# use a built-in model trained on tissuenet\n",
    "# (the first time you run this cell the model will download)\n",
    "model = models.CellposeModel(gpu=device, model_type='tissuenet')\n",
    "\n",
    "test_dataset = TissueNetDataset(root_dir='woodshole', split='test', prediction_type=prediction_type)\n",
    "test_loader = DataLoader(test_dataset, batch_size=1)\n",
    "\n",
    "### IMPORTANT: these are the channels used for the segmentation\n",
    "# the first one is the channel to segment, and the second one is the optional nuclear channel\n",
    "# red = 1\n",
    "# green = 2\n",
    "# blue = 3\n",
    "\n",
    "channels = [2, 1]\n",
    "\n",
    "masks_cp = []\n",
    "for idx, (image, mask) in enumerate(test_loader):\n",
    "    image = image.cpu().detach().numpy()\n",
    "    mask_cp, flows, styles = model.eval(image, diameter=25, channels=channels)\n",
    "    masks_cp.append(mask_cp)\n",
    "    \n",
    "    fig, axes = plt.subplots(1,4,figsize=(20, 20),sharex=True,sharey=True,squeeze=False)\n",
    "    \n",
    "    image = np.squeeze(image)\n",
    "    image = np.vstack((image, np.zeros_like(image)[:1]))\n",
    "    image = image.transpose(1,2,0)\n",
    "    \n",
    "    axes[0][0].imshow(image)\n",
    "    axes[0][0].title.set_text('Raw')\n",
    "    \n",
    "    axes[0][1].imshow(create_lut(mask_cp))\n",
    "    axes[0][1].title.set_text('Predicted Labels')\n",
    "\n",
    "    axes[0][2].imshow(flows[0])\n",
    "    axes[0][2].title.set_text('Predicted cellpose')\n",
    "    \n",
    "    axes[0][3].imshow(flows[2])\n",
    "    axes[0][3].title.set_text('Predicted cell probability')\n",
    "    \n",
    "    \n",
    "    if idx == 2:\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d710c00e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# we could also use the nuclear channel ONLY and run a nuclear model in cellpose\n",
    "# we set the second channel = 0 because we do not have an additional channel now\n",
    "channels = [1, 0]\n",
    "\n",
    "# initialize nuclei model (can also try \"cyto\" model if this doesn't work)\n",
    "# the \"nuclei\" model in cellpose has been trained on lots of nuclear data (but not the tissuenet dataset)\n",
    "# the \"cyto\" model in cellpose has been trained on many cellular images (but not the tissuenet dataset)\n",
    "model = models.CellposeModel(gpu=device, model_type='nuclei')\n",
    "\n",
    "masks_cp = []\n",
    "for idx, (image, mask) in enumerate(test_loader):\n",
    "    image = image.cpu().detach().numpy()\n",
    "    mask_cp, flows, styles = model.eval(image, diameter=20, channels=channels)\n",
    "    masks_cp.append(mask_cp)\n",
    "    \n",
    "    fig, axes = plt.subplots(1,4,figsize=(20, 20),sharex=True,sharey=True,squeeze=False)\n",
    "    \n",
    "    image = np.squeeze(image)\n",
    "    image = np.vstack((image, np.zeros_like(image)[:1]))\n",
    "    image = image.transpose(1,2,0)\n",
    "    \n",
    "    axes[0][0].imshow(image)\n",
    "    axes[0][0].title.set_text('Raw')\n",
    "    \n",
    "    axes[0][1].imshow(create_lut(mask_cp))\n",
    "    axes[0][1].title.set_text('Predicted Labels')\n",
    "\n",
    "    axes[0][2].imshow(flows[0])\n",
    "    axes[0][2].title.set_text('Predicted cellpose')\n",
    "    \n",
    "    axes[0][3].imshow(flows[2])\n",
    "    axes[0][3].title.set_text('Predicted cell probability')\n",
    "    \n",
    "    if idx == 2:\n",
    "        break"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4593600a",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "##### **Local Shape Descriptors**\n",
    "\n",
    "- Another example of auxiliary learning is [**LSDs**](https://localshapedescriptors.github.io/). This embedding encodes object shape similarly but is computed in a defined gaussian constrained to each label. This allows for consistent gradients regardless of object shapes which makes it a good candidate for segmentation of complex objects such as neurons in large electron microscopy datasets. \n",
    "\n",
    "\n",
    "- The LSDs are combined with nearest neighbor affinities to improve the boundary representations. The improved affinities then produce nice segmentations when using a hierarchical agglomeration approach and can be easily parallelized to allow for scaling to massive volumes. \n",
    "\n",
    "![example_image](static/lsd_schematic.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf570532-444e-411b-b452-7576f873d2d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# import lsds, calculate on a small patch and visualize the descriptor components\n",
    "\n",
    "from lsd.train import local_shape_descriptor\n",
    "\n",
    "file = random.choice(train_nuclei)\n",
    "\n",
    "nuclei = imread(file)[0:64, 0:64]\n",
    "raw = imread(file.replace('_nuclei_masks', ''))[:, 0:64, 0:64]\n",
    "\n",
    "#just to visualize\n",
    "raw = np.vstack((raw, np.zeros_like(raw)[:1]))\n",
    "raw = raw.transpose(1,2,0)\n",
    "\n",
    "lsds = local_shape_descriptor.get_local_shape_descriptors(\n",
    "              segmentation=nuclei,\n",
    "              sigma=(5,)*2,\n",
    "              voxel_size=(1,)*2)\n",
    "\n",
    "fig, axes = plt.subplots(\n",
    "            1,\n",
    "            6,\n",
    "            figsize=(20, 20),\n",
    "            sharex=False,\n",
    "            sharey=True,\n",
    "            squeeze=False)\n",
    "  \n",
    "axes[0][0].imshow(np.squeeze(lsds[0]), cmap='jet')\n",
    "axes[0][0].title.set_text('Mean offset Y')\n",
    "\n",
    "axes[0][1].imshow(np.squeeze(lsds[1]), cmap='jet')\n",
    "axes[0][1].title.set_text('Mean offset X')\n",
    "\n",
    "axes[0][2].imshow(np.squeeze(lsds[2]), cmap='jet')\n",
    "axes[0][2].title.set_text('Covariance Y-Y')\n",
    "\n",
    "axes[0][3].imshow(np.squeeze(lsds[3]), cmap='jet')\n",
    "axes[0][3].title.set_text('Covariance X-X')\n",
    "\n",
    "axes[0][4].imshow(np.squeeze(lsds[4]), cmap='jet')\n",
    "axes[0][4].title.set_text('Covariance Y-X')\n",
    "\n",
    "axes[0][5].imshow(np.squeeze(lsds[5]), cmap='jet')\n",
    "axes[0][5].title.set_text('Size')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "795f27af-1f0d-4d1e-a519-1ea4666793e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# slightly modify our dataset just for simplicity\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",
    "                 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.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",
    "    def __len__(self):\n",
    "        return len(self.raw_files)\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",
    "    def augment_data(self, raw, mask, padding):\n",
    "        \n",
    "        transform = A.Compose([\n",
    "              A.RandomCrop(\n",
    "                  width=self.crop_size,\n",
    "                  height=self.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\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 = raw.transpose([1,2,0])\n",
    "        \n",
    "        mask = np.expand_dims(mask, axis=0)\n",
    "        mask = mask.transpose([1,2,0])\n",
    "                \n",
    "        # just do this regardless of split to make val/test faster for demo purposes\n",
    "        padding = self.get_padding(self.crop_size, self.padding_size)\n",
    "        raw, mask = self.augment_data(raw, mask, padding)\n",
    "            \n",
    "        raw = raw.transpose([2,0,1])\n",
    "        mask = mask.transpose([2,0,1])\n",
    "        \n",
    "        mask, border = erode(\n",
    "                    mask[0],\n",
    "                    iterations=1,\n",
    "                    border_value=1)\n",
    "\n",
    "        affs = compute_affinities(mask, nhood=[[0,1],[1,0]])\n",
    "                        \n",
    "        lsds = local_shape_descriptor.get_local_shape_descriptors(\n",
    "              segmentation=mask,\n",
    "              sigma=(5,)*2,\n",
    "              voxel_size=(1,)*2)\n",
    "\n",
    "        lsds = lsds.astype(np.float32)\n",
    "        affs = affs.astype(np.float32)\n",
    "                                        \n",
    "        return raw, lsds, affs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0d178b0-b3a6-438c-ac67-5ba357a45a9e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# visualize batch\n",
    "\n",
    "train_dataset = TissueNetDataset(root_dir='woodshole', split='train', crop_size=64)\n",
    "\n",
    "raw, lsds, affs = 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(\n",
    "            1,\n",
    "            7,\n",
    "            figsize=(20, 20),\n",
    "            sharex=False,\n",
    "            sharey=True,\n",
    "            squeeze=False)\n",
    "  \n",
    "axes[0][0].imshow(np.squeeze(lsds[0]), cmap='jet')\n",
    "axes[0][0].title.set_text('Mean offset Y')\n",
    "\n",
    "axes[0][1].imshow(np.squeeze(lsds[1]), cmap='jet')\n",
    "axes[0][1].title.set_text('Mean offset X')\n",
    "\n",
    "axes[0][2].imshow(np.squeeze(lsds[2]), cmap='jet')\n",
    "axes[0][2].title.set_text('Covariance Y-Y')\n",
    "\n",
    "axes[0][3].imshow(np.squeeze(lsds[3]), cmap='jet')\n",
    "axes[0][3].title.set_text('Covariance X-X')\n",
    "\n",
    "axes[0][4].imshow(np.squeeze(lsds[4]), cmap='jet')\n",
    "axes[0][4].title.set_text('Covariance Y-X')\n",
    "\n",
    "axes[0][5].imshow(np.squeeze(lsds[5]), cmap='jet')\n",
    "axes[0][5].title.set_text('Size')\n",
    "\n",
    "axes[0][6].imshow(np.squeeze(affs[0]+affs[1]), cmap='jet')\n",
    "axes[0][6].title.set_text('Affs')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "855fe179-8ed2-45d3-b0b6-d03f761bffd4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# we need two output heads for our network, one for lsds and one for affinities\n",
    "# to do this we will subclass torch.nn.Module and create our UNet inside\n",
    "# before we had a single final convolution. Now we have one for each head.\n",
    "# then in the forward pass we pass our image through our unet and then the output through each head\n",
    "\n",
    "class MtlsdModel(torch.nn.Module):\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        in_channels,\n",
    "        num_fmaps,\n",
    "        fmap_inc_factors,\n",
    "        downsample_factors,\n",
    "        padding='same'\n",
    "    ):\n",
    "        super().__init__()\n",
    "\n",
    "        self.unet = UNet(\n",
    "            in_channels=in_channels,\n",
    "            num_fmaps=num_fmaps,\n",
    "            fmap_inc_factors=fmap_inc_factors,\n",
    "            downsample_factors=downsample_factors,\n",
    "            padding=padding)\n",
    "\n",
    "        self.lsd_head = torch.nn.Conv2d(in_channels=num_fmaps,out_channels=6, kernel_size=1)\n",
    "        self.aff_head = torch.nn.Conv2d(in_channels=num_fmaps,out_channels=2, kernel_size=1)\n",
    "\n",
    "    def forward(self, input):\n",
    "\n",
    "        z = self.unet(input)\n",
    "        lsds = self.lsd_head(z)\n",
    "        affs = self.aff_head(z)\n",
    "\n",
    "        return lsds, affs\n",
    "\n",
    "# We want to combine the lsds and affs losses and minimize the sum\n",
    "# we can do this by subclassing our loss function (torch.nn.MSELoss) and overriding the forward method\n",
    "\n",
    "class CombinedLoss(torch.nn.MSELoss):\n",
    "\n",
    "    def __init__(self):\n",
    "        super(CombinedLoss, self).__init__()\n",
    "\n",
    "    def forward(self, lsds_prediction, lsds_target, affs_prediction, affs_target):\n",
    "\n",
    "        loss1 = super(CombinedLoss, self).forward(lsds_prediction,lsds_target)\n",
    "        loss2 = super(CombinedLoss, self).forward(affs_prediction, affs_target)\n",
    "        \n",
    "        return loss1 + loss2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52eeeabc-09f4-4570-90fb-c4784d9c31bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(42)\n",
    "\n",
    "d_factors = [[2,2],[2,2],[2,2]]\n",
    "\n",
    "in_channels=2\n",
    "num_fmaps=32\n",
    "fmap_inc_factors=4\n",
    "\n",
    "net = MtlsdModel(in_channels,num_fmaps,fmap_inc_factors,d_factors)\n",
    "\n",
    "loss_fn = CombinedLoss().to(device)\n",
    "\n",
    "net = net.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "542a39b6-5c07-42bb-acde-10fcdea0a8c5",
   "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",
    "dtype = torch.FloatTensor\n",
    "\n",
    "# set optimizer\n",
    "learning_rate = 1e-4\n",
    "optimizer = torch.optim.Adam(net.parameters(), lr=learning_rate)\n",
    "\n",
    "# set activation\n",
    "activation = torch.nn.Sigmoid()\n",
    "\n",
    "### create datasets\n",
    "\n",
    "train_dataset = TissueNetDataset(root_dir='woodshole', split='train', crop_size=64)\n",
    "test_dataset = TissueNetDataset(root_dir='woodshole', split='test', crop_size=128)\n",
    "val_dataset = TissueNetDataset(root_dir='woodshole', split='test', val_split=True, crop_size=64)\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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c191c382-6301-46ae-84ec-cdaa54d11c19",
   "metadata": {},
   "outputs": [],
   "source": [
    "# update our training step to have two logits and two predictions\n",
    "\n",
    "def model_step(model, loss_fn, optimizer, feature, gt_lsds, gt_affs, activation, train_step=True):\n",
    "    \n",
    "    # zero gradients if training\n",
    "    if train_step:\n",
    "        optimizer.zero_grad()\n",
    "        \n",
    "    # forward\n",
    "    lsd_logits, affs_logits = model(feature)\n",
    "\n",
    "    loss_value = loss_fn(lsd_logits, gt_lsds, affs_logits, gt_affs)\n",
    "    \n",
    "    # backward if training mode\n",
    "    if train_step:\n",
    "        loss_value.backward()\n",
    "        optimizer.step()\n",
    "        \n",
    "    lsd_output = activation(lsd_logits)\n",
    "    affs_output = activation(affs_logits)\n",
    "   \n",
    "    outputs = {\n",
    "        'pred_lsds': lsd_output,\n",
    "        'pred_affs': affs_output,\n",
    "        'lsds_logits': lsd_logits,\n",
    "        'affs_logits': affs_logits,\n",
    "    }\n",
    "    \n",
    "    return loss_value, outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9f0fb35-41e6-4ee9-a550-45ba27fb143e",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# update our training loop to do both lsds and affs\n",
    "\n",
    "# set flags\n",
    "net.train() \n",
    "loss_fn.train()\n",
    "step = 0\n",
    "\n",
    "with tqdm(total=training_steps) 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, gt_lsds, gt_affs in tmp_loader:\n",
    "            gt_lsds = gt_lsds.to(device)\n",
    "            gt_affs = gt_affs.to(device)\n",
    "            feature = feature.to(device)\n",
    "                        \n",
    "            #print(label.shape, feature.shape)\n",
    "                    \n",
    "            loss_value, pred = model_step(net, loss_fn, optimizer, feature, gt_lsds, gt_affs, activation)\n",
    "            writer.add_scalar('loss',loss_value.cpu().detach().numpy(),step)\n",
    "            step += 1\n",
    "            pbar.update(1)\n",
    "            \n",
    "            if step % 100 == 0:\n",
    "                net.eval()\n",
    "                tmp_val_loader = iter(test_loader)\n",
    "                acc_loss = []\n",
    "                for feature, gt_lsds, gt_affs in tmp_val_loader:                    \n",
    "                    gt_lsds = gt_lsds.to(device)\n",
    "                    gt_affs = gt_affs.to(device)\n",
    "                    feature = feature.to(device)\n",
    "                    loss_value, _ = model_step(net, loss_fn, optimizer, feature, gt_lsds, gt_affs, activation, 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(np.mean(acc_loss))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7636afc0-7bed-4808-9c66-5dc52ba95796",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# visualize a few predictions - have the lsds helped to improve the affinities?\n",
    "# For a future challenge you could try using a weighted combined loss and watershed + agglomeration to get strong segmentations\n",
    "\n",
    "net.eval()\n",
    "\n",
    "activation = torch.nn.Sigmoid()\n",
    "\n",
    "for idx, (image, gt_lsds, gt_affs) in enumerate(test_loader):\n",
    "    image = image.to(device)\n",
    "    lsds_logits, affs_logits = net(image)\n",
    "    pred_lsds = activation(lsds_logits)\n",
    "    pred_affs = activation(affs_logits)\n",
    "        \n",
    "    image = np.squeeze(image.cpu())\n",
    "    gt_lsds = np.squeeze(gt_lsds.cpu().numpy())\n",
    "    gt_affs = np.squeeze(gt_affs.cpu().numpy())\n",
    "    \n",
    "    pred_lsds = np.squeeze(pred_lsds.cpu().detach().numpy())\n",
    "    pred_affs = np.squeeze(pred_affs.cpu().detach().numpy())\n",
    "    \n",
    "    fig, axes = plt.subplots(1,3,figsize=(20, 20),sharex=True,sharey=True,squeeze=False)\n",
    "    \n",
    "    image = np.vstack((image, np.zeros_like(image)[:1]))\n",
    "    image = image.transpose(1,2,0)\n",
    "    \n",
    "    axes[0][0].imshow(image)\n",
    "    axes[0][0].title.set_text('Raw')\n",
    "  \n",
    "    axes[0][1].imshow(np.squeeze(pred_lsds[0]), cmap='jet')\n",
    "    axes[0][1].imshow(np.squeeze(pred_lsds[1]), cmap='jet', alpha=0.5)\n",
    "    axes[0][1].title.set_text('Mean offsets')\n",
    "\n",
    "    axes[0][2].imshow(np.squeeze(pred_affs[0]+pred_affs[1]), cmap='jet')\n",
    "    axes[0][2].title.set_text('Affs')\n",
    "    \n",
    "    if idx == 2:\n",
    "        break"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3464d363-fcfc-4595-8f79-250c09bc9cc2",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "<hr style=\"height:2px;\"><div class=\"alert alert-block alert-success\"><h1>Checkpoint 4</h1>\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f63a89b",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "## Further learning\n",
    "\n",
    "* Instance segmentation can be challenging and this exercise just scratches the surface of what is possible.\n",
    "\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): http://funkey.science/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 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}