notebooks/DriftCorrection.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8dc41621",
   "metadata": {},
   "source": [
    "# NanoPyx \"Codeless\" Jupyter Notebook\n",
    "\n",
    "This notebook allows you to perform drift correction of 2D timelapses\n",
    "\n",
    "To use this notebook you don't need to interact with any code, just run cells in order and a graphical user interface will pop-up showcasing the parameters for each step.\n",
    "When running this notebook on Colab, it will automatically prompt you to connect to your google Drive, from where you can select the input data.\n",
    "To save the output of this notebook, make sure to select the \"save output\" option. Output will be saved on the same directory as the loaded image. If using an example dataset it will be saved on the path Python is currently running (if locally running the notebook) or on Colab files (if running on Colab, files button on the left bar).\n",
    "If you found this work useful for your research please consider citing our [preprint](https://www.biorxiv.org/content/10.1101/2023.08.13.553080v1)  and out Github repo [![DOI](https://zenodo.org/badge/505388398.svg)](https://zenodo.org/badge/latestdoi/505388398)\n",
    "\n",
    "**Drift Correction** Romain F Laine et al 2019 J. Phys. D: Appl. Phys. 52 163001\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4f7acce",
   "metadata": {
    "cellView": "form"
   },
   "outputs": [],
   "source": [
    "#@title Install NanoPyx, import necessary libraries and connect to Google Drive\n",
    "import sys\n",
    "IN_COLAB = 'google.colab' in sys.modules\n",
    "if IN_COLAB:\n",
    "    !pip install -q \"nanopyx[colab]\"\n",
    "    from google.colab import output\n",
    "    output.enable_custom_widget_manager()\n",
    "    from google.colab import drive\n",
    "    drive.mount('/content/drive')\n",
    "else:\n",
    "    !pip install -q \"nanopyx[jupyter]\"\n",
    "\n",
    "import io\n",
    "import os\n",
    "import cv2 as cv\n",
    "import skimage\n",
    "import nanopyx\n",
    "import stackview\n",
    "import numpy as np\n",
    "import tifffile as tiff\n",
    "import matplotlib as mpl\n",
    "import ipywidgets as widgets\n",
    "from PIL import Image\n",
    "from IPython.display import display, clear_output\n",
    "from matplotlib import pyplot as plt\n",
    "\n",
    "from nanopyx.core.utils.easy_gui import EasyGui\n",
    "from nanopyx.core.utils.find_files import find_files\n",
    "from nanopyx.data.download import ExampleDataManager\n",
    "\n",
    "cwd = os.getcwd()\n",
    "image_folder = \"datasets\"\n",
    "image_files = []\n",
    "EDM = ExampleDataManager()\n",
    "example_datasets = EDM.list_datasets()\n",
    "\n",
    "_path = os.path.join(\"..\", image_folder)\n",
    "if os.path.exists(_path):\n",
    "    image_files += find_files(_path, \".tif\")\n",
    "if os.path.exists(image_folder):\n",
    "    image_files += find_files(image_folder, \".tif\")\n",
    "image_files += [\"Example dataset: \"+dataset for dataset in example_datasets]\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23241f85",
   "metadata": {
    "cellView": "form"
   },
   "outputs": [],
   "source": [
    "#@title Load image stack\n",
    "# Create a GUI\n",
    "gui_data = EasyGui(\"Data Loader\")\n",
    "global own_data\n",
    "own_data = True\n",
    "\n",
    "def on_button_select_own(b):\n",
    "    clear_output()\n",
    "    gui_data.add_label(\"Select data to use:\")\n",
    "    gui_data.add_file_upload(\"upload\")\n",
    "    gui_data.add_dropdown(\"cmaps\", description=\"Colormap:\",\n",
    "                          options=sorted(list(mpl.colormaps)),\n",
    "                          value=\"viridis\", remember_value=True)\n",
    "    gui_data.add_button(\"load_data_own\", description=\"Load data\")\n",
    "    gui_data[\"load_data_own\"].on_click(on_button_load_data_clicked)\n",
    "    gui_data.show()\n",
    "\n",
    "def on_button_select_example(b):\n",
    "    clear_output()\n",
    "    gui_data.add_label(\"Select data to use:\")\n",
    "    gui_data.add_dropdown(\"data_source\", options=image_files,\n",
    "                    value=\"Example dataset: \"+example_datasets[4], remember_value=True)\n",
    "    gui_data.add_dropdown(\"cmaps\", description=\"Colormap:\",\n",
    "                          options=sorted(list(mpl.colormaps)),\n",
    "                          value=\"viridis\", remember_value=True)\n",
    "    gui_data.add_button(\"load_data\", description=\"Load data\")\n",
    "    gui_data[\"load_data\"].on_click(on_button_load_data_clicked_example)\n",
    "    gui_data.show()\n",
    "\n",
    "def on_button_load_data_clicked(b):\n",
    "    clear_output()\n",
    "    gui_data.show()\n",
    "    global dataset_original\n",
    "    global own_data\n",
    "    own_data = True\n",
    "    # disable button\n",
    "    gui_data[\"load_data_own\"].disabled = True\n",
    "    gui_data[\"load_data_own\"].description = \"Loading...\"\n",
    "    dataset_original = tiff.imread(gui_data[\"upload\"].selected)\n",
    "    gui_data[\"load_data_own\"].disabled = False\n",
    "    gui_data[\"load_data_own\"].description = \"Load data\"\n",
    "    gui_data._main_display.children = gui_data._main_display.children + (stackview.slice(dataset_original, colormap=gui_data[\"cmaps\"].value, continuous_update=True),)\n",
    "\n",
    "def on_button_load_data_clicked_example(b):\n",
    "    clear_output()\n",
    "    gui_data.show()\n",
    "    global dataset_original\n",
    "    global own_data\n",
    "    own_data = False\n",
    "    # disable button\n",
    "    gui_data[\"load_data\"].disabled = True\n",
    "    gui_data[\"load_data\"].description = \"Loading...\"\n",
    "\n",
    "    if gui_data[\"data_source\"].value.startswith(\"Example dataset: \"):\n",
    "        dataset_name = gui_data[\"data_source\"].value.replace(\n",
    "            \"Example dataset: \", \"\")\n",
    "        dataset_original = EDM.get_ZipTiffIterator(dataset_name, as_ndarray=True)\n",
    "        gui_data._main_display.children = gui_data._main_display.children + (stackview.slice(dataset_original, continuous_update=True, colormap=gui_data[\"cmaps\"].value),)\n",
    "    else:\n",
    "        dataset_original = skimage.io.imread(gui_data[\"data_source\"].value)\n",
    "        gui_data._main_display.children = gui_data._main_display.children + (stackview.slice(dataset_original, continuous_update=True, colormap=gui_data[\"cmaps\"].value),)\n",
    "\n",
    "    # enable button\n",
    "    gui_data[\"load_data\"].disabled = False\n",
    "    gui_data[\"load_data\"].description = \"Load data\"\n",
    "    gui_data.save_settings()\n",
    "\n",
    "gui_data.add_button(\"use_own_data\", description=\"Use Own data\")\n",
    "gui_data[\"use_own_data\"].on_click(on_button_select_own)\n",
    "gui_data.add_button(\"use_example_data\", description=\"Use Example data\")\n",
    "gui_data[\"use_example_data\"].on_click(on_button_select_example)\n",
    "gui_data.show()\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "448e86ee",
   "metadata": {},
   "source": [
    "# Drift Correction Parameters:  \n",
    "\n",
    "Requires an image stack with shape: (time, rows, columns)  \n",
    "\n",
    "- **Reference Frame:** Which frame to be used as reference. Either always use the first frame (better for fixed cells) or the previous frame (better for live cells).  \n",
    "- **Max Expected Drift:** Maximum amount of expected drift in pixels.  \n",
    "- **Time Averaging:** Whether to register each individual frame, if using 1, or how many frames to average before calculating drift correction (better for single molecule data). Output keeps the original number of frames, single frame drift estimation is calculated by interpolating using the calculated drift of the averaged image stack.\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9d7adb0",
   "metadata": {
    "cellView": "form"
   },
   "outputs": [],
   "source": [
    "#@title Create drift correction GUI\n",
    "from nanopyx.methods.drift_alignment.estimator import DriftEstimator\n",
    "\n",
    "gui_drift = EasyGui(\"Drift Correction\")\n",
    "\n",
    "def on_button_align(b):\n",
    "    clear_output()\n",
    "    gui_drift.show()\n",
    "    if gui_drift[\"ref\"].value == \"First frame\":\n",
    "        ref_option = 0\n",
    "    else:\n",
    "        ref_option = 1\n",
    "    avg = gui_drift[\"time_averaging\"].value\n",
    "    max_drift = gui_drift[\"max\"].value\n",
    "    global dataset_aligned\n",
    "    gui_drift[\"align\"].disabled = True\n",
    "    gui_drift[\"align\"].description = \"Aligning...\"\n",
    "    \n",
    "    estimator = DriftEstimator()\n",
    "    dataset_aligned = estimator.estimate(\n",
    "        dataset_original,\n",
    "        time_averaging=avg,\n",
    "        max_expected_drift=max_drift,\n",
    "        ref_option=ref_option,\n",
    "        apply=True)\n",
    "    \n",
    "    if gui_drift[\"save\"].value:\n",
    "        if own_data:\n",
    "            path = gui_data[\"upload\"].selected_path\n",
    "            name = gui_data[\"upload\"].selected_filename.split(\".\")[0]\n",
    "            drift_table_path = path + os.sep + name\n",
    "            save_path = path + os.sep + name + \"_aligned.tif\"\n",
    "        else:\n",
    "            name = gui_data[\"data_source\"].value.replace(\"Example dataset: \", \"\")\n",
    "            drift_table_path = path + os.sep + name\n",
    "            save_path = name + \"_aligned.tif\"\n",
    "        tiff.imwrite(save_path, dataset_aligned)\n",
    "        tmp = []\n",
    "        for key in estimator.estimator_table.params.keys():\n",
    "            tmp.append(((key, estimator.estimator_table.params[key])))\n",
    "        tmp.append(estimator.estimator_table.drift_table)\n",
    "        np.save(drift_table_path + \"_drift_table.npy\", tmp)\n",
    "    gui_drift[\"align\"].disabled = False\n",
    "    gui_drift[\"align\"].description = \"Align\"\n",
    "    gui_drift._main_display.children = gui_drift._main_display.children + (stackview.slice(dataset_aligned, colormap=gui_drift[\"cmaps\"].value, continuous_update=True),)\n",
    "\n",
    "gui_drift.add_label(\"Drift Correction parameters:\")\n",
    "gui_drift.add_dropdown(\"ref\", description=\"Reference frame\", options=[\"First frame\", \"Previous frame\"], value=\"First frame\")\n",
    "gui_drift.add_int_slider(\"max\", description=\"Max expected drift\", min=0, max=1000, value=10)\n",
    "gui_drift.add_int_slider(\"time_averaging\", description=\"Time averaging\", min=1, max=dataset_original.shape[0], value=1)\n",
    "gui_drift.add_dropdown(\"cmaps\", description=\"Colormap:\",\n",
    "                      options=sorted(list(mpl.colormaps)),\n",
    "                      value=\"viridis\", remember_value=True)\n",
    "gui_drift.add_checkbox(\"save\", description=\"Save Output\", value=True)\n",
    "gui_drift.add_button(\"align\", description=\"Align\")\n",
    "gui_drift[\"align\"].on_click(on_button_align)\n",
    "gui_drift.show()\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c8afefe",
   "metadata": {},
   "source": [
    "## Use the following cell only if you have a previously calculated drift table\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55eb1110",
   "metadata": {
    "cellView": "form"
   },
   "outputs": [],
   "source": [
    "#@title Create apply drift correction GUI\n",
    "import os\n",
    "import sys\n",
    "\n",
    "NANOPYX_INSTALLED = \"nanopyx\" in sys.modules\n",
    "if not NANOPYX_INSTALLED:\n",
    "    IN_COLAB = \"google.colab\" in sys.modules\n",
    "    if IN_COLAB:\n",
    "        !pip install -q \"nanopyx[colab]\"\n",
    "        from google.colab import output\n",
    "        output.enable_custom_widget_manager()\n",
    "        from google.colab import drive\n",
    "        drive.mount('/content/drive')\n",
    "    else:\n",
    "        !pip install -q \"nanopyx[jupyter]\"\n",
    "\n",
    "import stackview\n",
    "import tifffile as tiff\n",
    "import matplotlib as mpl\n",
    "from nanopyx.core.utils.easy_gui import EasyGui\n",
    "from nanopyx.methods.drift_alignment.corrector import DriftCorrector\n",
    "from IPython.display import display, clear_output\n",
    "\n",
    "gui_drift_apply = EasyGui(\"Apply Drift Correction\")\n",
    "\n",
    "def on_button_apply(b):\n",
    "    clear_output()\n",
    "    gui_drift_apply.show()\n",
    "\n",
    "    gui_drift_apply[\"Align image\"].disabled = True\n",
    "    gui_drift_apply[\"Align image\"].description = \"Aligning...\"\n",
    "\n",
    "    drift_table_path = gui_drift_apply[\"upload\"].selected\n",
    "    img_path = gui_drift_apply[\"upload image\"].selected\n",
    "\n",
    "    corrector = DriftCorrector()\n",
    "    corrector.load_estimator_table(drift_table_path)\n",
    "    aligned_image = corrector.apply_correction(tiff.imread(img_path))\n",
    "\n",
    "    if gui_drift_apply[\"save\"].value:\n",
    "        path = gui_drift_apply[\"upload image\"].selected_path\n",
    "        name = gui_drift_apply[\"upload image\"].selected_filename.split(\".\")[0]\n",
    "        save_path = path + os.sep + name + \"_aligned.tif\"\n",
    "        tiff.imwrite(save_path, aligned_image)\n",
    "\n",
    "\n",
    "\n",
    "    gui_drift_apply[\"Align image\"].disabled = False\n",
    "    gui_drift_apply[\"Align image\"].description = \"Align\"\n",
    "    gui_drift_apply._main_display.children = gui_drift_apply._main_display.children + (stackview.slice(aligned_image, colormap=gui_drift_apply[\"cmaps\"].value, continuous_update=True),)\n",
    "\n",
    "gui_drift_apply.add_label(\"Load drift table:\")\n",
    "gui_drift_apply.add_file_upload(\"upload\")\n",
    "gui_drift_apply.add_label(\"Load image to align:\")\n",
    "gui_drift_apply.add_file_upload(\"upload image\")\n",
    "gui_drift_apply.add_dropdown(\"cmaps\", description=\"Colormap:\",\n",
    "                      options=sorted(list(mpl.colormaps)),\n",
    "                      value=\"viridis\", remember_value=True)\n",
    "gui_drift_apply.add_checkbox(\"save\", description=\"Save Output\", value=True)\n",
    "gui_drift_apply.add_button(\"Align image\", description=\"Align image\")\n",
    "gui_drift_apply[\"Align image\"].on_click(on_button_apply)\n",
    "gui_drift_apply.show()\n",
    "\n"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 5
}