using_siam_unet.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Siamese U-Net Quickstart"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**IMPORTANT**: Two packages packages need to be installed manually before running bio-image-unet: CUDA and PyTorch. To install CUDA 11.1 which is officially supported by PyTorch, navigate to [its installation page](https://developer.nvidia.com/cuda-11.1.1-download-archive) and follow the instructions onscreen. Because PyTorch depends on your CUDA installation version, it will need to be installed manually as well, through [the official PyTorch website](https://pytorch.org/get-started/locally/). Select the correct distribution of CUDA on this webpage and run the command in your terminal. bio-image-unet doesn't depend on a specific version of CUDA and has been tested with PyTorch 1.7.0+.\n",
    "\n",
    "## 1. Introduction\n",
    "\n",
    "The Siamese U-Net is an improvement on the original U-Net architecture. It adds an additional additional encoder that encodes an additional frame other than the frame that we are trying to predict, and uses a cross correlation layer to make the best inference for the data, detailed in [this paper](https://pubmed.ncbi.nlm.nih.gov/31927473/). This repository contains an implementation of this network.\n",
    "\n",
    "This document serves both as a \"Getting Started\" tutorial and \"Best Practices\" documentation. If you need help using a class (one that is directly under the `biu.siam_unet` directory), the examples in this notebook should be helpful. If you need help using a function or if you are interested in seeing the full list of parameters for a certain function, you can run `help(whichever_interesting_function)` in the interactive shell or take a peek at the source code. \n",
    "\n",
    "Finally, to import the Siamese U-Net package, write `import biu.siam_unet as unet`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Data preparation\n",
    "\n",
    "Because Siam UNet requires an additional input for training, we need to utilize an additional frame and use the appropriate dataloader for that. For the purpose of this notebook, I will call the frame which we are trying to infer \"current frame\", and the frame which is before the current frame the \"previous frame.\" "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### If your input image is not a movie"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from bio_image_unet.siam_unet.helpers.generate_siam_unet_input_imgs import generate_coupled_image_from_self\n",
    "from pathlib import Path\n",
    "import os\n",
    "\n",
    "# specify where the training data for vanilla u-net is located\n",
    "training_data_loc = '/home/longyuxi/Documents/mount/deeptissue_training/training_data/amnioserosa/yokogawa/image'\n",
    "training_data_loc = Path(training_data_loc)\n",
    "\n",
    "# create a separate folder for storing Siam-UNet input images\n",
    "siam_training_data_loc = training_data_loc.parent / \"siam_image\"\n",
    "siam_training_data_loc.mkdir(exist_ok=True)\n",
    "\n",
    "### multiprocessing accelerated, equivalent to \n",
    "## for img in training_data_loc.glob('*.tif'):\n",
    "##     generate_coupled_image_from_self(str(img), str(siam_training_data_loc / img.name))\n",
    "\n",
    "import multiprocessing\n",
    "\n",
    "imglist = training_data_loc.glob('*.tif')\n",
    "def handle_image(img):\n",
    "    generate_coupled_image_from_self(str(img), str(siam_training_data_loc / img.name))\n",
    "\n",
    "p = multiprocessing.Pool(10)\n",
    "_ = p.map(handle_image, imglist)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tifffile\n",
    "a = tifffile.imread('/home/longyuxi/Documents/mount/deeptissue_training/training_data/leading_edge/eCad/image/00.tif')\n",
    "print(a.shape)\n",
    "\n",
    "generate_coupled_image_from_self('/home/longyuxi/Documents/mount/deeptissue_training/training_data/leading_edge/eCad/image/00.tif', 'temp.tif')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### If you know which frame you drew the label with"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The dataloader in `siam_unet_cosh` takes an image that results from concatenating the previous frame with the current frame. If you already know which frame of which movie you want to train on, you can create this concatenated data using `generate_siam_unet_input_imgs.py`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "movie_dir = '/media/longyuxi/H is for HUGE/docmount backup/unet_pytorch/training_data/test_data/new_microscope/21B11-shgGFP-kin-18-bro4.tif' # change this\n",
    "frame = 10 # change this\n",
    "out_dir = './training_data/training_data/yokogawa/siam_data/image/' # change this\n",
    "\n",
    "\n",
    "\n",
    "from bio_image_unet.siam_unet.helpers.generate_siam_unet_input_imgs import generate_coupled_image\n",
    "\n",
    "generate_coupled_image(movie_dir, frame, out_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### If you don't know which frame you drew the label with"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you have frames and labels, but you don't know which frame of which movie each frame comes from, you can use  `find_frame_of_image`. This function takes your query and compares it against a list of tif files you specify through the parameter `search_space`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "image_name = f'./training_data/training_data/yokogawa/lateral_epidermis/image/83.tif'\n",
    "\n",
    "razer_local_search_dir = '/media/longyuxi/H is for HUGE/docmount backup/all_movies'\n",
    "tifs_names = ['21B11-shgGFP-kin-18-bro4', '21B25_shgGFP_kin_1_Pos0', '21C04_shgGFP_kin_2_Pos4', '21C26_shgGFP_Pos12', '21D16_shgGFPkin_Pos7']\n",
    "search_space = [razer_local_search_dir + '/' + t + '.tif' for t in tifs_names]\n",
    "\n",
    "from bio_image_unet.siam_unet.helpers.find_frame_of_image import find_frame_of_image\n",
    "\n",
    "find_frame_of_image(image_name, search_space=search_space)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This function not only outputs what it finds to stdout, but also creates a machine readable output, location of which specified by `machine_readable_output_filename`, about which frames it is highly confident with at locating (i.e. an MSE of < 1000 and matching frame numbers). This output can further be used by `generate_siam_unet_input_images.py`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from bio_image_unet.siam_unet.helpers.generate_siam_unet_input_imgs import utilize_search_result\n",
    "\n",
    "utilize_search_result(f'./training_data/training_data/yokogawa/amnioserosa/search_result_mr.txt', f'./training_data/test_data/new_microscope', f'./training_data/training_data/yokogawa/amnioserosa/label/', f'./training_data/training_data/yokogawa/siam_amnioserosa_sanitize_test/')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, organize the labels and images in a way similar to this shown. An example can be found at `training_data/lateral_epidermis/yokogawa_siam-u-net`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```\n",
    "training_data/lateral_epidermis/yokogawa_siam-u-net\n",
    "|\n",
    "├── image\n",
    "│   ├── 105.tif\n",
    "│   ├── 111.tif\n",
    "│   ├── 120.tif\n",
    "│   ├── 121.tif\n",
    "│   ├── 1.tif\n",
    "│   ├── 2.tif\n",
    "│   ├── 3.tif\n",
    "│   ├── 5.tif\n",
    "│   ├── 7.tif\n",
    "│   └── 83.tif\n",
    "└── label\n",
    "    ├── 105.tif\n",
    "    ├── 111.tif\n",
    "    ├── 120.tif\n",
    "    ├── 121.tif\n",
    "    ├── 1.tif\n",
    "    ├── 2.tif\n",
    "    ├── 3.tif\n",
    "    ├── 5.tif\n",
    "    ├── 7.tif\n",
    "    └── 83.tif\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Training is simple. For example:\n",
    "\n",
    "Note that the *invert* flag must be set to True for DeepTissue to segment the cells correctly (that is, the prediction would have white cell boundaries and black background)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from bio_image_unet.siam_unet import *\n",
    "\n",
    "dataset = 'amnioserosa/old_scope'\n",
    "base_dir = '/home/longyuxi/Documents/mount/deeptissue_training/training_data/'\n",
    "\n",
    "# path to training data (images and labels with identical names in separate folders)\n",
    "dir_images = f'{base_dir}/{dataset}/siam_image/'\n",
    "dir_masks = f'{base_dir}/{dataset}/label/'\n",
    "\n",
    "print('starting to create training dataset')\n",
    "print(f'dir_images: {dir_images}')\n",
    "print(f'dir_masks: {dir_masks}')\n",
    "# create training data set\n",
    "data = DataProcess([dir_images, dir_masks], data_path='../delete_this_data', dilate_mask=0, aug_factor=10, create=True, clip_threshold=(0.2, 99.8), dim_out=(256, 256), shiftscalerotate=(0, 0, 0))\n",
    "\n",
    "save_dir = f'/home/longyuxi/Documents/mount/trained_networks_new_siam/siam/{dataset}'\n",
    "# create trainer\n",
    "training = Trainer(data ,num_epochs=500 ,batch_size=12, load_weights=False, lr=0.0001, n_filter=32, save_iter=False, save_dir=save_dir, loss_function='logcoshTversky', loss_params=(1.5, 0.6))\n",
    "\n",
    "\n",
    "training.start()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note here that the value of the `n_filter` parameter is set to `32`. This is the depth of the first convolution layer in the U-Net architecture. Assignment of this value should depend on the amount of RAM that one's GPU has, but `32` layers has worked well for us and runs on a GTX 1080 with 8GB RAM. Should you change this value, you need to change the `n_filter` parameter when running `Predict` (next section) as well."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A note on loss functions\n",
    "\n",
    "One can switch the loss functions used to train the network. There are three options for the `loss_function` parameter: `BCEDice`, `Tversky` and `logcoshTversky`, and the two-element tuple `loss_params` to accompany each loss function.\n",
    "\n",
    "#### BCE Dice Loss\n",
    "\n",
    "The BCE Dice loss is a combination of the binary cross entropy loss (BCELoss) and Dice loss. Both are commonly used loss functions for image segmentation, and the BCE Dice loss is a weighted sum of them. The BCE Dice loss is defined as:\n",
    "\n",
    "$$\n",
    "\\mathcal{L_\\text{BCE Dice}} = \\alpha * \\mathcal{L_\\text{BCE}} + \\beta * \\mathcal{L_\\text{Dice}}\n",
    "$$\n",
    "\n",
    "We implemented $\\mathcal{L_\\text{BCE}}$, $\\mathcal{L_\\text{Dice}}$ and $\\mathcal{L_\\text{BCE Dice}}$ based on https://github.com/achaiah/pywick/blob/master/pywick/losses.py.\n",
    "\n",
    "#### Tversky Loss\n",
    "\n",
    "Tversky loss is defined as follows (as described in https://arxiv.org/pdf/1706.05721.pdf):\n",
    "\n",
    "$$\n",
    "\\mathcal{L_\\text{Tversky}} = \\frac{\\text{TP} + S}{\\text{TP} + \\alpha * \\text{FP} + \\beta * \\text{FN} + S}\n",
    "$$\n",
    "\n",
    "Where TP, FP, and FN indicate true positives, false positives and false negatives respectively. $S$ is a smooth factor that is set to 1 by default and can't be changed in the current implementation. $\\alpha$ and $\\beta$ are weights on how much importance FP and FN should have, respectively. A higher $\\alpha$ would make the network more sensitive on false positives, and are passed in through `loss_params`.\n",
    "\n",
    "#### logcoshTversky Loss\n",
    "\n",
    "As suggested by its name, the logcoshTversky loss is defined as \n",
    "\n",
    "$$\n",
    "\\mathcal{L_\\text{Tversky}} = \\log{(\\cosh{(\\mathcal{L_\\text{Tversky}})})}\n",
    "$$\n",
    "\n",
    "And the parameters of `loss_params` are passed into $\\alpha$ and $\\beta$ of the inner Tversky loss.\n",
    "\n",
    "#### Our experience with these loss functions:\n",
    "\n",
    "`BCEDice` with `loss_params=(1.0, 1.0)` has consistently produced decent segmentation for us, so we think `BCEDice` can be used as a \"default\" loss function for running a dataset for the first time. `Tversky` and `logcoshTversky` losses, when trained with good $\\alpha$ and $\\beta$ parameters, can produce even better training results. The optimal weight of $\\alpha$ and $\\beta$ for the Tversky loss functions seems to depend on proportion of each image that is label vs background, since they alter the weight of false positives and false negatives."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Predict"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Predicting is simple as well. Just swap in the parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load package\n",
    "from bio_image_unet.siam_unet import *\n",
    "import os\n",
    "os.nice(10)\n",
    "from  bio_image_unet.siam_unet.helpers import tif_to_mp4\n",
    "\n",
    "base_dir = './'\n",
    "out_dir = f'{base_dir}/predicted_out'\n",
    "model = f'{base_dir}/models/siam_bce_amnio/model_epoch_100.pt'\n",
    "\n",
    "tif_file = f'{base_dir}/training_data/test_data/new_microscope/21C04_shgGFP_kin_2_Pos4.tif'\n",
    "\n",
    "result_file = f'{out_dir}/siam_bce_amnio_100_epochs_21C04_shgGFP_kin_2_Pos4.tif'\n",
    "out_mp4_file = result_file[:-4] + '.mp4'\n",
    "\n",
    "print('starting to predict file')\n",
    "# predict file \n",
    "predict = Predict(tif_file, result_file, model, invert=False, resize_dim=(512, 512))\n",
    "# convert to mp4\n",
    "tif_to_mp4.convert_to_mp4(result_file, output_file=out_mp4_file, normalize_to_0_255=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Additionally, to evaluate the model's performance with different losses, one can also train the model across different models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "For each image in the training dataset, run siam unet to predict.\n",
    "\"\"\"\n",
    "\n",
    "from pathlib import *\n",
    "\n",
    "from bio_image_unet.siam_unet import *\n",
    "import glob\n",
    "import logging\n",
    "\n",
    "def predict_all_training_data(image_folder_prefix, model_folder_prefix, model_loss_functions, datasets, output_directory):\n",
    "    image_folder_prefix = Path(image_folder_prefix)\n",
    "    model_folder_prefix = Path(model_folder_prefix)\n",
    "    datasets = [Path(d) for d in datasets]\n",
    "    output_directory = Path(output_directory)\n",
    "    for dataset in datasets:\n",
    "        for model_loss_function in model_loss_functions:\n",
    "            try:\n",
    "                current_model = Path(model_folder_prefix / model_loss_function / dataset / 'model.pt')\n",
    "                for image in glob.glob((str) (image_folder_prefix / dataset) + \"/image/*.tif\"):\n",
    "                    image_name = image.split('/')[-1]\n",
    "                    result_name = Path(output_directory / dataset / Path(image_name[:-4] + '_' + model_loss_function + '.tif'))\n",
    "                    _ = Predict(image, result_name, current_model, invert=False)\n",
    "                    # _ = Predict(image, result_name, current_model, invert=False, resize_dim=None, n_filter=32)\n",
    "            except:\n",
    "                logging.error('{} in {} failed to execute'.format(model_loss_function, dataset))\n",
    "                \n",
    "\n",
    "if __name__ == '__main__':\n",
    "    # BEGIN Full dataset\n",
    "    folders = [\"amnioserosa/yokogawa\", \"lateral_epidermis/40x\", \"lateral_epidermis/60x\", \"lateral_epidermis/yokogawa\", \"leading_edge/eCad\", \"leading_edge/myosin\", \"leading_edge/yokogawa_eCad\", \"nodes/old_scope\", \"nodes/yokogawa\"]\n",
    "    model_loss_functions = ['siam_bce_dice','siam_logcoshtversky', 'siam_tversky', 'siam_logcoshtversky_08_02', 'siam_logcoshtversky_15_06', 'siam_logcoshtversky_02_08', \"siam_logcoshtversky_06_15\", 'siam_tversky_08_02', 'siam_tversky_15_06']\n",
    "    # END Full dataset\n",
    "\n",
    "    # BEGIN Toy dataset\n",
    "    # folders = [\"lateral_epidermis/40x\"]\n",
    "    # model_loss_functions = ['siam_bce_dice','siam_logcoshtversky', 'siam_tversky', 'siam_logcoshtversky_08_02', 'siam_logcoshtversky_15_06']\n",
    "    # END Toy dataset\n",
    "\n",
    "    predict_all_training_data(image_folder_prefix='/home/longyuxi/Documents/mount/deeptissue_training/training_data', model_loss_functions=model_loss_functions, model_folder_prefix='/home/longyuxi/Documents/mount/trained_networks', datasets=folders, output_directory='/home/longyuxi/Documents/mount/deeptissue_test/output_new_shape')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Appendix: An annotated structure of the siam_unet package"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Below is an annotated structure of the siam_unet package. Use `help(function)` to read the docstring of each function for a better understanding.\n",
    "\n",
    "```\n",
    "Package                                         Use\n",
    "\n",
    ".\n",
    "├── __init__.py\n",
    "├── data.py                                     dataloader script\n",
    "├── siam_unet.py                                Siam U-Net model\n",
    "├── train.py                                    training script\n",
    "├── losses.py                                   loss functions\n",
    "├── predict.py                                  prediction script\n",
    "├── helpers                                     helper functions (usually not \n",
    "                                                        so useful except the \n",
    "                                                        ones mentioned in this notebook)\n",
    "                                                        \n",
    "│   ├── average_tifs.py                             averages a list of tiff files\n",
    "│   ├── create_pixel_value_histogram.py             creates histograms for the \n",
    "                                                        pixel values in tif \n",
    "                                                        files. Useful for \n",
    "                                                        debugging during training\n",
    "│   ├── cuda_test.py                                tests cuda functionality\n",
    "│   ├── extract_frame_of_movie.py                   extract a certain frame of a \n",
    "                                                        tif movie \n",
    "│   ├── find_frame_of_image.py                      finds the frame number of \n",
    "                                                        a given query image \n",
    "                                                        within search_space.\n",
    "│   ├── generate_plain_image.py                     generates a plain image\n",
    "│   ├── generate_siam_unet_input_imgs.py            generates a coupled image \n",
    "                                                        for Siam U-Net training\n",
    "│   ├── low_mem_tif_utils.py                        utilities for handling tif \n",
    "                                                        files with low memory \n",
    "                                                        usage\n",
    "│   ├── threshold_images.py                         thresholds each frame of a \n",
    "                                                        tif movie\n",
    "│   ├── tif_to_mp4.py                               uses ffmpeg to convert a tif \n",
    "                                                        movie to mp4\n",
    "│   └── util.py                                     various utilities. see docstring\n",
    "\n",
    "```"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "6bbe29e2308ccd271c4ed360a09e65ae469839e520c3df2570e8ecd9e52abb09"
  },
  "kernelspec": {
   "display_name": "Python 3.8.5 64-bit ('base': conda)",
   "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}