04_bonus_Noise2Noise/n2n.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Noise2Noise\n",
    "\n",
    "CARE network you trained in the first regression exercises require that you acquire pairs\n",
    "of high and low SNR. However, this often not possible. One such case is when it is simply\n",
    "not possible to acquire high SNR images.\n",
    "\n",
    "What to do when you are stuck with just noisy images? We also have seen Noise2Void, which\n",
    "is a self-supervised method that can be trained on noisy images. But there are other \n",
    "supervised approaches that can be trained on noisy images only, such as Noise2Noise. \n",
    "\n",
    "Noise2Noise relies on the same assumption than Noise2Void: the noise is pixel-independent.\n",
    "Therefore, if you supervise your network to guess a noisy image from another one, the network\n",
    "will converge to a denoised image. Of course, this only works if the two noisy images are\n",
    "very similar.\n",
    "\n",
    "To acquire data for Noise2Noise, one can simply image the same region of interest twice!\n",
    "Indeed, pixel-independent noise (as opposed to structured noise) will be completely independent\n",
    "between neighboring pixels and as well as between the two noisy images.\n",
    "\n",
    "In this notebook, we will again use the [Careamics](https://careamics.github.io) library.\n",
    "\n",
    "## Reference\n",
    "\n",
    "Lehtinen, Jaakko, et al. \"[Noise2Noise: Learning image restoration without clean data.](https://arxiv.org/abs/1803.04189)\" arXiv preprint arXiv:1803.04189 (2018).\n",
    "\n",
    "\n",
    "<div class=\"alert alert-block alert-success\"><h3>Objectives</h3>\n",
    "    \n",
    "- Understand the differences between CARE, Noise2Noise and Noise2Void\n",
    "- Train Noise2Noise with CAREamics\n",
    "  \n",
    "</div>\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import shutil\n",
    "from pathlib import Path\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import tifffile\n",
    "from careamics_portfolio import PortfolioManager\n",
    "\n",
    "from careamics import CAREamist\n",
    "from careamics.config import create_n2n_configuration"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<hr style=\"height:2px;\">\n",
    "\n",
    "## Part 1: Prepare the data\n",
    "\n",
    "The N2N SEM dataset consists of EM images with 7 different levels of noise:\n",
    "\n",
    "- Image 0 is recorded with 0.2 us scan time\n",
    "- Image 1 is recorded with 0.5 us scan time\n",
    "- Image 2 is recorded with 1 us scan time\n",
    "- Image 3 is recorded with 1 us scan time\n",
    "- Image 4 is recorded with 2.1 us scan time\n",
    "- Image 5 is recorded with 5.0 us scan time\n",
    "- Image 6 is recorded with 5.0 us scan time and is the avg. of 4 images\n",
    "\n",
    "Let's have a look at them.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download files\n",
    "portfolio = PortfolioManager()\n",
    "root_path = Path(\"./data\")\n",
    "files = portfolio.denoising.N2N_SEM.download(root_path)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualize training data\n",
    "\n",
    "In this cell we can see the different levels of noise in the SEM dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load images\n",
    "train_image = tifffile.imread(root_path / \"denoising-N2N_SEM.unzip/SEM/train.tif\")\n",
    "print(f\"Train image shape: {train_image.shape}\")\n",
    "\n",
    "# plot image\n",
    "fig, ax = plt.subplots(1, 2, figsize=(10, 10))\n",
    "ax[0].imshow(train_image[0,100:356, 500:756], cmap=\"gray\")\n",
    "ax[0].set_title(\"Train image highest noise level\")\n",
    "ax[1].imshow(train_image[-1, 100:356, 500:756], cmap=\"gray\")\n",
    "ax[1].set_title(\"Train image lowest noise level\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-info\"><h3>Task 1: Explore the data</h3>\n",
    "\n",
    "Vizualize each different noise level!\n",
    "\n",
    "</div>\n",
    "\n",
    "<hr style=\"height:2px;\">\n",
    "\n",
    "## Part 2: Create the configuraion\n",
    "\n",
    "As in the Noise2Void exercise, a good CAREamics pipeline starts with a configuration!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "training_config = create_n2n_configuration(\n",
    "    experiment_name=\"N2N_SEM\",\n",
    "    data_type=\"array\",\n",
    "    axes=\"SYX\",\n",
    "    patch_size=[128, 128],\n",
    "    batch_size=128,\n",
    "    num_epochs=50,\n",
    "    logger=\"tensorboard\"\n",
    ")\n",
    "\n",
    "# Visualize training configuration\n",
    "print(training_config)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<hr style=\"height:2px;\">\n",
    "\n",
    "## Part 3: Train the network\n",
    "\n",
    "In this part, we create our training engine (`CAREamics`) and start training the network."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create the engine\n",
    "careamist = CAREamist(source=training_config)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-info\"><h3>Task 2: Which data to choose?</h3>\n",
    "\n",
    "How would you train a network to denoise images of 1 us scan time? Which images do you think could be used as input and which as target?\n",
    "\n",
    "Set the `train_source` and `train_target` accordingly and train the network.\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the training data and targets\n",
    "train_data = train_image[[2, 2, 2, 2, 2, 3, 3, 3, 3, 3], ...]\n",
    "train_target = train_image[[0, 1, 3, 4, 5, 0, 1, 3, 4, 5], ...]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "careamist.train(\n",
    "    train_source=train_data,\n",
    "    train_target=train_target\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-success\"><h1>Checkpoint 1: Training N2N</h1>\n",
    "</div>\n",
    "\n",
    "\n",
    "<hr style=\"height:2px;\">\n",
    "\n",
    "## Part 4: Prediction"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's load the test data and predict on it to assess how well the network performs!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load images\n",
    "test_image = tifffile.imread(root_path / \"denoising-N2N_SEM.unzip/SEM/test.tif\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prediction = careamist.predict(source=test_image[2], tile_size=(256, 256), axes=\"YX\", tta_transforms=False)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualize predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 2, figsize=(10, 10))\n",
    "ax[0].imshow(test_image[-1], cmap=\"gray\")\n",
    "ax[0].set_title(\"Test image lowest noise level\")\n",
    "ax[1].imshow(prediction, cmap=\"gray\")\n",
    "ax[1].set_title(\"Prediction\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fi, ax = plt.subplots(1, 2, figsize=(15, 15))\n",
    "vim  = test_image[0].min()\n",
    "vmax = test_image[0].max()\n",
    "ax[0].imshow((prediction.squeeze())[1000:1128, 500:628], cmap=\"gray\",vmin=vim, vmax=vmax)\n",
    "ax[0].set_title(\"Prediction\")\n",
    "ax[1].imshow(test_image[-1].squeeze()[1000:1128, 500:628], cmap=\"gray\", vmin=vim, vmax=vmax)\n",
    "ax[1].set_title(\"Test image lowest noise level\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-info\"><h3>Task 3: Different noise pairs</h3>\n",
    "\n",
    "Can you further improve your results by usign different `source` and `target`?\n",
    "\n",
    "How would you train a network to denoise all images, rather than just the 1 us ones?\n",
    "\n",
    "Try it and be creative!\n",
    "\n",
    "</div>"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "cmcs_l",
   "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.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}