Examples/use_unet/ensemble_mode_worm4-clear.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3DeeCellTracker Demo: ensemble mode\n",
    "\n",
    "This notebook shows how to use 3DeeCellTracker to track cells in ensemble mode.\n",
    "The demo data can be found in the \"ensemble_mode_worm4\" folder downloaded from https://osf.io/dt76c/\n",
    "\n",
    "**The basic procedures:**\n",
    "- A. Import packages\n",
    "- B. Initialize the parameters for tracking\n",
    "- C. Prepare images to be tracked, and the pre-trained U-Net and FFN models.\n",
    "- D. Optimize segmentation parameters and segment the image at volume 1.\n",
    "- E. Correct the segmentation at volume 1 (in other image processing software)\n",
    "- F. Optimize tracking parameters.\n",
    "- G. Track following volumes.\n",
    "\n",
    "\n",
    "**Please run following codes according to the instructions**\n",
    "- Some codes can be skipped if unnecessary."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A. Import packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T00:45:54.266854Z",
     "start_time": "2021-05-28T00:45:54.229081Z"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "import os\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "\n",
    "from IPython.core.display import display, HTML\n",
    "from CellTracker.unet3d import unet3_b\n",
    "from matplotlib.animation import FuncAnimation, ArtistAnimation\n",
    "from CellTracker.tracker import Tracker\n",
    "display(HTML(\"<style>.container { width:95% !important; }</style>\"))\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## B. Initialize the parameters for tracking\n",
    "###  Set parameters for the tracker\n",
    "**Image parameters**\n",
    "- volume: number of the volumes (time points) of the images to be tracked\n",
    "- siz_xyz: size of each 3D image: (height, width, depth), unit: voxels\n",
    "- z_xy_ratio: resolution (um/voxels) ratio between z (depth) and x-y plane\n",
    "- z_scaling: (integer), the factor for interpolating the unbalanced z_xy resolution\n",
    "- miss_frame: (list), the volume number with serious problems and should be skipped during tracking. default value is an empty list: []\n",
    "\n",
    "**Segmentation parameters**\n",
    "- noise_level: a value close to the averaged non-cell region intensity, which helps the program to ignore the background noises.\n",
    "- min_size: the possible minimal cell size, which helps the program to ignore small non-cell regions.\n",
    "- Note: To change them after this initialization, please use \".set_segmentation()\" which will delete the cached segmentation files\n",
    "\n",
    "**Tracking parameters**\n",
    "- beta_tk: set it higher/lower to get more coherent/independent predictions for cell positions.\n",
    "- lambda_tk: set it higher/lower to get more coherent/independent predictions for cell positions.\n",
    "- maxiter_tk: (integer) the number of iteration for the tracking algorithm (FFN+PR-GLS), the higher, the more accurate (but slower)\n",
    "- ensemble: (integer or bool) in ensemble mode, it should be an positive interger, indicating the maximum number of predictions to be averaged. The default value is \"False\", indicating single mode.\n",
    "\n",
    "**Paths**\n",
    "- folder_path: the path of the folder to store the data, model, and results.\n",
    "    - \"./xxx\" indicates a folder with name \"xxx\" under the directory containing this jupyter notebook.\n",
    "- image_name: the names of the images to be tracked. \n",
    "    - \"aligned_t%03i_z%03i.tif\" indicates file names like: \"aligned_t002_z011.tif\", \"aligned_t502_z101.tif\", etc.\n",
    "- unet_model_file: the name of the pre-trained unet model file\n",
    "- ffn_model_file: the name of the pre-trained ffn model file\n",
    "\n",
    "### Notes: \n",
    "**Multiple folders were automatically created to store data, model, and results**\n",
    "- /data: to store the images to be tracked\n",
    "- /auto_vol1: to store the automatically generated segmentation results at volume 1\n",
    "- /manual_vol1: to store the manually corrected segmentation results at volume 1\n",
    "- /tracking_information: to store the coordinates information of the tracked cells\n",
    "- /models: to store the pretrained and retrained models\n",
    "- /unet: to store the cached predictions of cell regions\n",
    "- /track_results_EnsembleDstrbtMode: to store the tracking results (images of labels)\n",
    "- /anim: to store the animation of the tracking process (used for diagnosis)\n",
    "- /models/unet_weights: to store the retrained unet weights\n",
    "- Note: in case one folder have been created, the program will not overwrite it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T00:47:06.924231Z",
     "start_time": "2021-05-28T00:47:06.889428Z"
    }
   },
   "outputs": [],
   "source": [
    "tracker = Tracker(\n",
    "    volume_num=80, siz_xyz=(168, 401, 128), z_xy_ratio=1, z_scaling=1, miss_frame=[79],\n",
    "    noise_level=200, min_size=50, beta_tk=1000, lambda_tk=0.00001, maxiter_tk=10, ensemble=20,\n",
    "    folder_path=os.path.abspath(\"./worm4\"), image_name=\"raw_t%04i_z%04i.tif\", \n",
    "    unet_model_file=\"unet3_pretrained_worm4.h5\",ffn_model_file=\"ffn_pretrained.h5\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## C. Prepare images and pre-trained models\n",
    "**Prepare images**\n",
    "- Move images to the created \"/data\" folder. The names of the images should be consistent with the parameter \"image_name\"\n",
    "- Move the pre-trained 3D U-Net and FFN model files to the created \"/models\" folder"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## D. Optimize segmentation parameters and segment the image at volume 1.\n",
    "\n",
    "**D1. Modify the segmentation parameters (optional)**\n",
    "- Please modify these segmentation parameters only when the following segmentation results are unsatisfied\n",
    "- Note: the cached segmentation in \"/unet\" folder will be deleted after this modification"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T00:46:44.385865Z",
     "start_time": "2021-05-28T00:46:44.351315Z"
    }
   },
   "outputs": [],
   "source": [
    "tracker.set_segmentation(noise_level=200, min_size=50)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**D2. Segment cells at volume 1**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T01:03:30.321638Z",
     "start_time": "2021-05-28T01:02:25.509171Z"
    }
   },
   "outputs": [],
   "source": [
    "tracker.load_unet()\n",
    "tracker.segment_vol1()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**D3. Draw the results of segmentation (Max projection)**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T01:12:56.504696Z",
     "start_time": "2021-05-28T01:12:31.176496Z"
    }
   },
   "outputs": [],
   "source": [
    "anim_seg = tracker.draw_segresult(percentile_high=99.8)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**D4. Show segmentation in each layer**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T01:13:01.221093Z",
     "start_time": "2021-05-28T01:13:01.106899Z"
    }
   },
   "outputs": [],
   "source": [
    "HTML(anim_seg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## E. Manually correct the segmentation at volume 1 and load it.\n",
    "**E1. Manual correction and move files to the folder**\n",
    "- Delete non-cell regions;\n",
    "- Add cells not labeled; \n",
    "- correct cell boundaries in other software (such as ITK-SNAP)\n",
    "\n",
    "- save the corrected segmentation as 2D images (with different z)\n",
    "- Move the corrected segmentation to the created \"/manual_vol1\" folder"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**E2. Load the manually corrected segmentation**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T01:15:54.027005Z",
     "start_time": "2021-05-28T01:15:53.527729Z"
    }
   },
   "outputs": [],
   "source": [
    "tracker.load_manual_seg()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**E3. Re-train the U-Net using the manual segmentation (optional)**\n",
    "- Please do this only if the prediction by U-Net is poor (e.g. many cells were not ignored or many non-cell regions were identified as cells)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T01:41:59.630956Z",
     "start_time": "2021-05-28T01:16:09.813035Z"
    },
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "tracker.retrain_unet()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Select an step with satisfied segmentation\n",
    "- If the initial segmentation is better, set step=0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T02:06:05.752753Z",
     "start_time": "2021-05-28T02:04:37.885200Z"
    }
   },
   "outputs": [],
   "source": [
    "tracker.select_unet_weights(step=10)\n",
    "tracker.set_segmentation(del_cache=True)\n",
    "tracker.segment_vol1()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**E4. Interpolate cells to make more accurate/smooth cell boundary**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T02:02:32.324888Z",
     "start_time": "2021-05-28T02:02:27.551990Z"
    }
   },
   "outputs": [],
   "source": [
    "tracker.interpolate_seg()\n",
    "tracker.draw_manual_seg1()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**E5. Initiate variables required for tracking**\n",
    "- cal_subregions(): required by \"accurate correction\" in following procedures\n",
    "- check_multicells(): if multiple separated cells were labeled as the same one, generate an error message."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T02:07:40.041282Z",
     "start_time": "2021-05-28T02:07:36.233259Z"
    }
   },
   "outputs": [],
   "source": [
    "tracker.cal_subregions()\n",
    "tracker.load_ffn()\n",
    "tracker.initiate_tracking()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## F. Optimize tracking parameters. \n",
    "**F1. Modify tracking parameters if the test result is not satisfied (optional)**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T02:10:37.096752Z",
     "start_time": "2021-05-28T02:10:37.059173Z"
    }
   },
   "outputs": [],
   "source": [
    "tracker.set_tracking(beta_tk=1000, lambda_tk=0.00001, maxiter_tk=10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-14T04:54:59.219773Z",
     "start_time": "2021-05-14T04:54:59.151351Z"
    }
   },
   "source": [
    "**F2. Test a matching between volume 1 and a target volume, and show the FFN + PR-GLS process by an animation (5 iterations)**\n",
    "- Red circles: manually corrected cells in vol 1\n",
    "- Blue cross: cells detected by 3D U-Net in target volume\n",
    "- Blue arrows: predicted transformation (cell movements)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T02:11:51.354831Z",
     "start_time": "2021-05-28T02:10:41.588093Z"
    },
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "anim_tracking, results = tracker.match(target_volume=21)\n",
    "HTML(anim_tracking)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**F3. Show the accurate correction after the FFN + PR-GLS transformation**\n",
    "- Arrows show the correction of the cell positions\n",
    "- Cells without arrow: no correction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T02:12:14.007555Z",
     "start_time": "2021-05-28T02:12:12.904857Z"
    }
   },
   "outputs": [],
   "source": [
    "tracker.draw_correction(*results[2:])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**F4. Show the superimposed cells + labels before/after tracking**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T02:12:20.921021Z",
     "start_time": "2021-05-28T02:12:18.671159Z"
    },
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "tracker.draw_overlapping(*results[:3])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## G. Tracking following volumes\n",
    "**G1. Track and show the processes**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T04:02:22.761978Z",
     "start_time": "2021-05-28T02:12:46.838632Z"
    },
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "%matplotlib notebook\n",
    "fig, ax = tracker.subplots_tracking()\n",
    "tracker.track(fig, ax, from_volume=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**G2. Show the processes as an animation (for diagnosis)**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T04:11:33.933350Z",
     "start_time": "2021-05-28T04:11:27.159333Z"
    },
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "track_anim = tracker.replay_track_animation()\n",
    "HTML(track_anim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "position": {
    "height": "521.183px",
    "left": "897px",
    "right": "20px",
    "top": "146px",
    "width": "539px"
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}