resources/notebooks/N2V_denoising.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "94a07788-92e9-4e1f-a45f-4ac9a188cc57",
   "metadata": {},
   "source": [
    "### 1. Import dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0802461b-a09c-41ad-9a35-448c39100aad",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "import tifffile\n",
    "from tqdm.auto import tqdm\n",
    "import numpy as np\n",
    "from matplotlib import pyplot as pl\n",
    "\n",
    "from n2v.models import N2VConfig, N2V\n",
    "from n2v.utils.n2v_utils import manipulate_val_data\n",
    "from n2v.internals.N2V_DataGenerator import N2V_DataGenerator"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d9c012b-9313-404b-8e6e-fa743eab279a",
   "metadata": {},
   "source": [
    "### 2. Define main parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce9bcfd4-10bc-4d9f-8a60-fec4532ffa45",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Folder containing noisy images\n",
    "image_dir = './train_videos/'\n",
    "\n",
    "# Name of the resulting model. If you already have the model with this name, it will be overwritten\n",
    "model_name = 'wound_healing_n2v'\n",
    "\n",
    "# Path where the model will be saved once trained\n",
    "model_path = './models/'\n",
    "\n",
    "# How many epochs (rounds) the network will be trained. \n",
    "# Preliminary results can already be observed after a few (10-30) epochs, but a full training should run for 100-200 epochs. \n",
    "number_of_epochs = 50\n",
    "\n",
    "# N2V divides image into patches during training, this value is the size of the resulting patches\n",
    "patch_shape=(64, 64)\n",
    "\n",
    "# Input the percentage of your training dataset you want to use to validate the network during the training\n",
    "percentage_validation = 10\n",
    "\n",
    "# Limit the max number of generated patches in case of lower RAM\n",
    "max_patches = 10000"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c52ce2d-7c2b-4b41-8142-84ed8b9481db",
   "metadata": {},
   "source": [
    "### 3. Create data generator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ebfa42b4-373c-4620-9708-a5658225c426",
   "metadata": {},
   "outputs": [],
   "source": [
    "datagen = N2V_DataGenerator()\n",
    "images = datagen.load_imgs_from_directory(directory=image_dir, filter='*.tif', dims='TCYX')\n",
    "patches = datagen.generate_patches_from_list(images, shape=patch_shape)\n",
    "patches = patches[:max_patches]\n",
    "\n",
    "total_number_of_patches = len(patches)\n",
    "num_training_patches = total_number_of_patches - int(total_number_of_patches * (percentage_validation / 100))\n",
    "X = patches[:num_training_patches]\n",
    "X_val = patches[num_training_patches:]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7548e860-8550-4355-bf1e-2e294dde2a44",
   "metadata": {},
   "source": [
    "### 4. Define model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4258599-ce09-4e91-a6d5-da65a9b24014",
   "metadata": {},
   "outputs": [],
   "source": [
    "config = N2VConfig(\n",
    "    X, \n",
    "    blurpool=True,\n",
    "    skip_skipone=True,\n",
    "    unet_residual=False,\n",
    "    n2v_manipulator='median', \n",
    "    unet_kern_size=3, \n",
    "    unet_n_first=64,      \n",
    "    unet_n_depth=3, \n",
    "    train_steps_per_epoch=int(X.shape[0]/128), \n",
    "    train_epochs=number_of_epochs, \n",
    "    train_loss='mse', \n",
    "    batch_norm=True, \n",
    "    train_batch_size=128, \n",
    "    n2v_perc_pix=0.198, \n",
    "    n2v_patch_shape=patch_shape, \n",
    "    n2v_neighborhood_radius=5, \n",
    "    single_net_per_channel=False\n",
    ")\n",
    "\n",
    "model = N2V(config=config, name=model_name, basedir=model_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38a1cedf-ccf2-4c76-be9b-53a1050568a2",
   "metadata": {},
   "source": [
    "### 5. Train and save model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "767f894d-ad28-4d15-8fc6-6c4264883b3d",
   "metadata": {},
   "outputs": [],
   "source": [
    "history = model.train(X, X_val)\n",
    "\n",
    "model.export_TF(\n",
    "    name='N2V', \n",
    "    description=\"\", \n",
    "    authors=[\"\"],\n",
    "    test_img=X[0], axes='YXC',\n",
    "    patch_shape=patch_shape\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08197de5-68c4-48d4-9db6-fb9c7f30574c",
   "metadata": {},
   "source": [
    "### 6. Plot training and validation loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2995cca2-08b9-4b8c-b058-f039b364fcca",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from csbdeep.utils import plot_history\n",
    "\n",
    "plt.figure(figsize=(16, 5))\n",
    "plot_history(history, ['loss', 'val_loss']);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "212840d6-13c5-4215-bd17-5133e025077f",
   "metadata": {},
   "source": [
    "### 7. Prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56069f86-807d-4e1e-9069-2fb3e2551ffb",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_folder = \"./all_videos/\" # path to the data you want to predict\n",
    "\n",
    "data_filter = \"*.tif\" \n",
    "\n",
    "output_postfix = \"denoised\" # results will be saved in the same folder, with this postfix added to the name\n",
    "\n",
    "# If you want to predict with a different model,\n",
    "# uncomment this lines and add a path to the model:\n",
    "\n",
    "# model_name = \"custom model name\"\n",
    "# model_path = \"custom model path\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d2026d4f-9766-4cc8-afdc-8be8c04daed5",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = N2V(\n",
    "    config=None, \n",
    "    name=model_name, \n",
    "    basedir=model_path\n",
    ")\n",
    "\n",
    "samples = sorted(Path(data_folder).glob(data_filter))\n",
    "\n",
    "if len(samples) == 0:\n",
    "    print(f\"No data found in folder: {data_folder}\")\n",
    "\n",
    "for sample in tqdm(samples):\n",
    "    image = tifffile.imread(sample)\n",
    "    image = np.swapaxes(image, 1, -1)  # CYX -> YXC\n",
    "    \n",
    "    result = []\n",
    "    for timepoint in image:\n",
    "        pred = model.predict(timepoint, axes='YXC')\n",
    "        result.append(pred)\n",
    "        \n",
    "    result = np.stack(result)\n",
    "    result = np.swapaxes(result, -1, 1)\n",
    "\n",
    "    name = sample.stem\n",
    "    result_path = sample.parent / f\"{sample.stem}_{output_postfix}.tif\"\n",
    "    tifffile.imwrite(result_path, result, imagej=True, metadata={'axes': 'TCYX', 'mode': 'rgb'})"
   ]
  }
 ],
 "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.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}