Examples/use_unet/3D_U_Net_training-clear.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "697b0d44",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# 3DeeCellTracker Demo: Train 3D U-Net\n",
    "This notebook shows how to use 3DeeCellTracker to train a 3D U-Net.\n",
    "The demo data can be found in the \"unet_training\" folder downloaded from https://osf.io/dt76c/\n",
    "\n",
    "The basic procedures:\n",
    "\n",
    "- A. Import packages\n",
    "- B. Initialize the parameters for training\n",
    "- C. Load the train / validation datasets\n",
    "- D. Preprocess the datasets\n",
    "- E. Train the 3D U-Net\n",
    "- F. Select the best weights and save the model\n",
    "\n",
    "Please run following codes according to the instructions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ecb96b7",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## A. Import packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "limited-documentation",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T08:55:42.591834Z",
     "start_time": "2021-05-28T08:55:35.453228Z"
    }
   },
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "import os\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "from CellTracker.unet3d import TrainingUNet3D, unet3_a\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61d33354",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## B. Initialize the parameters for training\n",
    "- noise_level: This value should be roughly the intensity of the background pixels in the raw image. It is used in\n",
    "normalization to enhance the cells with weak intensity while ignore the background noises. If the two images (train and\n",
    "validation) have different noise level, choose a value somehow between them.\n",
    "\n",
    "- folder_path: This path is used to create a folder (if not exist) to store the data and model files. We recommend users\n",
    "to set it as “./xxxx” to create a folder named xxxx under the same directory containing the notebook file.\n",
    "\n",
    "- model: This should be a predefined 3D U-Net model. The simplest way is to use the default value unet3_a(). Advanced\n",
    "users can select other predefined model such as unet3_b(), unet3_c(), or a model defined by themselves (need to modify\n",
    "the “Import packages” cell to import the model)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96253465",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T08:55:53.931353Z",
     "start_time": "2021-05-28T08:55:43.525252Z"
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "trainer = TrainingUNet3D(noise_level=100, folder_path=os.path.abspath(\"./unet_01\"), model=unet3_a())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53622a5c",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## C. Load the train / validation datasets\n",
    "\n",
    "After running B, the program automatically generated several folders under the folder_path. Users should prepare the\n",
    "training data and validation data (see examples in the demo data) and store them into the “train_image”(raw 2D image\n",
    "sequence for training data), “train_label”(2D annotation of cell/non-cell regions for training data), “valid_image” and\n",
    "“valid_label” (for validation data), respectively.\n",
    "\n",
    "\n",
    "Run the code cell. The program will load and draw the images/annotations of the training and validation dataset by max-projection."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3952e893",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "trainer.load_dataset()\n",
    "trainer.draw_dataset()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c890ddc",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## D. Preprocess the datasets\n",
    "**D1. Show the normalized images (max projection)**\n",
    "\n",
    "Normalize the images and show the normalized images.\n",
    "If the normalization looks poor, e.g. the background look too bright, or the intensity of weak cells are not enhanced.\n",
    "Please go back to step B to modify the parameter “noise_level” and run all codes again."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20df5b6b",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "trainer.preprocess()\n",
    "trainer.draw_norm_dataset()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "329cab21",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "**D2. Show the divided images (part)**\n",
    "\n",
    "Divide the images into multiple sub-images (used as the input of 3D U-Net) and show a part of these sub-images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79516b0e",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "trainer.draw_divided_train_data()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e2560a95",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## E. Train the 3D U-Net\n",
    "Start to train the 3D U-Net. During the training, the program will draw the updated prediction of cell regions if the\n",
    "loss on validation data is reduced.\n",
    "\n",
    "By default, the program will train for 100 epochs. Users can manually stop the training by pressing Ctrl+C if the\n",
    "val_loss no longer decreases anymore."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c82b3bef",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4782a367",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## F. Select the best weights and save the model\n",
    "After training finished or manually stopped. Users should choose the best step which generate the best prediction of\n",
    "cell regions. Usually this should be the one with the lowest val_loss, but uses can select other step instead.\n",
    "\n",
    "The program will store the model with the chosen weights into the “models” folder with the name “unet3_pretrained.h5”"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48196a91",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "trainer.select_weights(step=97)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6447d3a1",
   "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()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}