notebooks/1_Single_Channel_Eval.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b3a8744b",
   "metadata": {},
   "source": [
    "# Part 1: Explore data, test and tune processing parameters."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "741c3fdd",
   "metadata": {},
   "source": [
    "### In the following notebook you will prepare for processing your images; test illumination correction parameters, stitching accuracy, and deconvolution."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "91b73c1b",
   "metadata": {},
   "source": [
    "## 1. Import packages. \n",
    "### *This must be done every time the notebook is started or restarted."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "278d48ba",
   "metadata": {},
   "source": [
    "Run cells using Ctrl+Enter or by pressing the play (triangle) button."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "211fcc99",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os as os\n",
    "from tqdm.notebook import tqdm\n",
    "import tkinter as tk\n",
    "from tkinter import filedialog\n",
    "from tkinter import simpledialog\n",
    "import pandas as pd\n",
    "from Kstitch.stitching import stitch_images\n",
    "from glob import glob\n",
    "from skimage.io.collection import alphanumeric_key\n",
    "import numpy as np\n",
    "from skimage.io import imread \n",
    "from skimage.io import imsave\n",
    "from skimage import io\n",
    "from skimage import exposure\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from itertools import chain, repeat\n",
    "import subprocess\n",
    "from datetime import datetime\n",
    "import imagej, scyjava\n",
    "import platform\n",
    "import stackview\n",
    "import Kview2\n",
    "from typing import List\n",
    "from skimage.transform import resize as skresize\n",
    "from scipy.fftpack import dct, idct\n",
    "import pyopencl as cl\n",
    "os_system = platform.system()\n",
    "current_dateTime = datetime.now()\n",
    "%matplotlib widget"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff5009c1",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "## 2. Define directory paths.  \n",
    "### *This must be done every time the notebook is started or restarted."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1083961",
   "metadata": {},
   "source": [
    "- If you are familiar with defining directories, enter them below.  If not, the other cells will assist with the necessary definitions. \n",
    "\n",
    "- If you have established the directories with Method B, from now on use Method A by copying the directories in their respective variable names."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0edad7a",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "### 2.1a Method A"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "1039737b",
   "metadata": {},
   "outputs": [],
   "source": [
    "base_dir = \"C:\\\\Users\\\\smith6jt\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "be45f614-5f53-4099-81d0-be16ece5577c",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Image folder is C:\\Users\\smith6jt\\KINTSUGI\\data\\1904_CC2B_raw.\n",
      "Stitching folder is C:\\Users\\smith6jt\\KINTSUGI\\data\\1904_CC2B_BaSiC_Stitched.\n",
      "Meta folder is C:\\Users\\smith6jt\\KINTSUGI\\data\\1904_CC2B_meta.\n"
     ]
    }
   ],
   "source": [
    "image_dir = os.path.join(base_dir, 'KINTSUGI', 'data', '1904_CC2B_raw')\n",
    "stitch_dir = image_dir.replace(\"_raw\", \"_BaSiC_Stitched\")\n",
    "meta_dir = stitch_dir.replace(\"_BaSiC_Stitched\", \"_meta\")\n",
    "project_file = os.path.join(meta_dir, \"project_data.txt\")\n",
    "print(f\"Image folder is {image_dir}.\")\n",
    "print(f\"Stitching folder is {stitch_dir}.\")\n",
    "print(f\"Meta folder is {meta_dir}.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e669516f",
   "metadata": {},
   "source": [
    "### 2.1b Method B"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3526acea",
   "metadata": {},
   "source": [
    "- First, assess the folder and file names.  For example, in CODEX output, the original folder name is something like CX_20-008_SP_CC2-B. We can shorten that now to make life easier down the road.  Shorten the folder name to remove redundant information.  In this example we choose \"2008_CC2B_raw\" making sure the '_raw' is the last part.\n",
    "\n",
    "- Running the following cell will bring up a dialog window where you will enter the new folder name and select the folder containing your images.  In the folder you select, make sure there are only a series of folders containing the images for each cycle.  \n",
    "\n",
    "- This folder will be renamed and designated as \"image_dir\" for the project.  A stitching output folder and a folder for metadata files will also be created.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "835ff0b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "root = tk.Tk()\n",
    "root.withdraw()\n",
    "\n",
    "user_choice = simpledialog.askstring(\"Shortened name\", \"Enter shortened file name with _raw at end\")\n",
    "image_dir_long = filedialog.askdirectory()\n",
    "head_dir, tail_dir = os.path.split(image_dir_long)\n",
    "base_dir, _ = os.path.split(head_dir)\n",
    "\n",
    "try:  \n",
    "    os.rename(image_dir_long, os.path.join(head_dir, user_choice))\n",
    "except FileNotFoundError:\n",
    "    print(\"The file or directory does not exist.\")\n",
    "except PermissionError:\n",
    "    print(\"you don't have permissions to rename the file\")\n",
    "except OSError as error:\n",
    "    print(f\"Error: {error}\")\n",
    "\n",
    "if os.path.isdir(os.path.join(head_dir, user_choice)) == True:\n",
    "    image_dir = os.path.join(head_dir, user_choice)\n",
    "    print(f\"{image_dir_long} renamed to {image_dir}\")\n",
    "\n",
    "else:\n",
    "    print(\"Error in renaming.\")\n",
    "    \n",
    "stitch_dir = image_dir.replace('_raw', '_BaSiC_Stitched')\n",
    "meta_dir = image_dir.replace('_raw', '_meta')\n",
    "os.makedirs(stitch_dir, exist_ok=True)\n",
    "os.makedirs(meta_dir, exist_ok=True)\n",
    "\n",
    "if os.path.isdir(stitch_dir) == True:\n",
    "    print(f\"Stitching output folder: {stitch_dir} created successfully.\")\n",
    "\n",
    "else:\n",
    "    print(\"Stitching output folder not created.\")\n",
    "\n",
    "if os.path.isdir(meta_dir) == True:\n",
    "    print(f\"Metadata folder: {meta_dir} created successfully.  Move metadata files to metadata folder.\")\n",
    "\n",
    "else:\n",
    "    print(\"Metadata folder not created.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "69f85fd4",
   "metadata": {},
   "source": [
    "### 2.2 Shorten cycle folder names"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aefe29ff",
   "metadata": {},
   "source": [
    "NOTE: Running this cell permanently alters the cycle folder names.  This only needs to be done once.\n",
    "\n",
    "- If the cycle folders have long names starting with the cycle numbers followed by a split character (e.g. the underscore in 'cyc001_reg001_200210_170925') this cell will shorten the folder name to everything before the first split character.\n",
    "\n",
    "- Enter the split character in quotations.  It is usually an underscore.\n",
    "\n",
    "- Note that the rest of the folder name will be deleted unless you add additional arguments to the os.rename function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39720611",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Rename long source folder names\n",
    "\n",
    "split_character = simpledialog.askstring(\"Split character\", \"Enter split character\")\n",
    "print(f'Split character is {split_character}.')\n",
    "if split_character == 'None':\n",
    "    print('No input given.')\n",
    "\n",
    "else:\n",
    "    for cyc_folder in os.listdir(image_dir):\n",
    "        new_name = cyc_folder.split(split_character)\n",
    "        cycle_dir = os.path.join(image_dir, cyc_folder)\n",
    "        cycle_dir_short = os.path.join(image_dir, new_name[0])\n",
    "        os.rename(cycle_dir, cycle_dir_short)\n",
    "        if os.path.isdir(cycle_dir_short) == True:\n",
    "            print(f\"{cycle_dir} renamed to {cycle_dir_short}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f774b4d",
   "metadata": {},
   "source": [
    "### 2.4 Import test set"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aae9d09c",
   "metadata": {},
   "source": [
    "- Processing very large image datasets takes considerable time.  To avoid having to repeat the process and to use the most efficient processing methods, the rest of this notebook allows for the discovery and tuning of parameters for each step.  \n",
    "\n",
    "- First choose and inspect a single z-plane of a single cycle/channel combination, and check the quality of illumination correction and stitching.  Then you will work with an entire stack to assess deconvolution and extended depth of focus (EDF).  \n",
    "\n",
    "- Starting with the nuclear staining channel or other channels necessary for segmentation is a good idea.\n",
    "\n",
    "- Make sure to assess channels with differing staining patterns and intensities.\n",
    "\n",
    "- Enter the z-plane, channel number, and pattern of the image_file name below.  The wildcard characters \"??\" are entered where the image tile numbers are so that all tiles are loaded.\n",
    "\n",
    "- The cycle folder names are assumed to have been shortened and in the format \"cyc00x\".  Running the cell will prompt to enter the cycle folder of interest or you may comment out the first line and uncomment the second where your folder is defined.\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d6948ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "# cycle_folder = filedialog.askdirectory()\n",
    "cycle_folder = os.path.join(image_dir, \"cyc003\")\n",
    "\n",
    "zplane = 8 \n",
    "channel = 1\n",
    "\n",
    "# image_file name below is derived from 1_000tt_Z0zz_CHc.tif where tt is the tile number with one leading zero, zz is the z-position with one leading zero, and c is the channel number.\n",
    "image_file = f'1_000??_Z0{str(zplane).zfill(2)}_CH{channel}.tif'\n",
    "print(f\"Using {cycle_folder}/{image_file}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e771989-98f8-4d57-93ed-1a1ccb934243",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "im_raw = sorted(glob(os.path.join(cycle_folder, image_file)), key = alphanumeric_key)\n",
    "im = io.imread_collection(im_raw)\n",
    "im_array_init = np.asarray(im)\n",
    "print(f\"Image shape is {im_array_init.shape} and type is {im_array_init.dtype}\\nMin value is {im_array_init.min()} and max value is {im_array_init.max()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7cd0407b",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "## 3. Illumination Correction"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b749cf08",
   "metadata": {},
   "source": [
    "### 3.1 Downsize and resize. \n",
    "\n",
    "Define the function to resize the images to a smaller and square dimension for estimating the uneven illumination.  Usually no changes need to be made here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0445640d",
   "metadata": {},
   "outputs": [],
   "source": [
    "RESIZE_ORDER = 1\n",
    "RESIZE_MODE = \"symmetric\"\n",
    "PRESERVE_RANGE = True\n",
    "OUTPUT_IMAGE = \"OutputImage\"\n",
    "\n",
    "def _resize_images_list(images_list: List, side_size: float = None, x_side_size: float = None, y_side_size: float = None):\n",
    "    if side_size is not None:\n",
    "        y_side_size = x_side_size = side_size\n",
    "    resized_images_list = []\n",
    "    for i, im in enumerate(images_list):\n",
    "        if im.shape[0] != x_side_size or im.shape[1] != y_side_size:\n",
    "            resized_images_list.append(skresize(\n",
    "                im, \n",
    "                (x_side_size, y_side_size), \n",
    "                order = RESIZE_ORDER, \n",
    "                mode = RESIZE_MODE,\n",
    "                preserve_range = PRESERVE_RANGE\n",
    "                )\n",
    "            )\n",
    "        else:\n",
    "            resized_images_list.append(im)\n",
    "    return resized_images_list"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8532a167",
   "metadata": {},
   "source": [
    "Run the resize function and print out some image information.  Usually no changes need to be made here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6403d6a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "_saved_size = im[0].shape\n",
    "working_size = 128\n",
    "nrows = ncols = working_size\n",
    "\n",
    "downsized_image = np.dstack(_resize_images_list(images_list=im, side_size=working_size))\n",
    "print(f\"Downsized image shape is {downsized_image.shape} and type is {downsized_image.dtype}\\nMin value is {downsized_image.min()} and max value is {downsized_image.max()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16dd88aa",
   "metadata": {},
   "source": [
    "- Plot the original image and histogram next to the resized image and histogram.  \n",
    "- Enter a number for \"i\" that is within the range of the number of images in your dataset for one channel.  \n",
    "- For this and the rest of the image/histogram cells, you can adjust the size of the plot by entering values for \"figsize.\"  \n",
    "- You can also change the colors of the images by changing \"cmap.\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "179a4e24",
   "metadata": {},
   "outputs": [],
   "source": [
    "i = 40\n",
    "\n",
    "hist, hist_centers = exposure.histogram(im_array_init[i])\n",
    "hist2, hist_centers2 = exposure.histogram(downsized_image[:,:,i])\n",
    "fig, axes = plt.subplots(1, 4, figsize=(12, 3))\n",
    "axes[0].imshow(im_array_init[i], cmap='turbo')\n",
    "axes[0].set_title('Original Image')\n",
    "axes[0].set_axis_off()\n",
    "axes[1].plot(hist_centers, hist, lw=2)\n",
    "axes[1].set_title('Histogram of Pixel Intensities')\n",
    "axes[2].imshow(downsized_image[:,:,i], cmap='turbo')\n",
    "axes[2].set_title('Downsized Image')\n",
    "axes[2].set_axis_off()\n",
    "axes[3].plot(hist_centers2, hist2, lw=2)\n",
    "axes[3].set_title('Histogram of Pixel Intensities')\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a58b10c5",
   "metadata": {},
   "source": [
    "### 3.2 Take Mean and Transform\n",
    "This cell takes the average of all the image tiles and standardizes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71d729d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "mean_image = np.mean(downsized_image, axis=2)\n",
    "mean_div = mean_image/np.mean(mean_image)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c8ff0e4",
   "metadata": {},
   "source": [
    "View the mean and standardized images as above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8c346dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "i = 40\n",
    "\n",
    "hist, hist_centers = exposure.histogram(mean_image)\n",
    "hist2, hist_centers2 = exposure.histogram(mean_div)\n",
    "fig, axes = plt.subplots(1, 4, figsize=(12, 3))\n",
    "axes[0].imshow(mean_image, cmap='turbo')\n",
    "axes[0].set_title('Downsized image')\n",
    "axes[0].set_axis_off()\n",
    "axes[1].plot(hist_centers, hist, lw=2)\n",
    "axes[1].set_title('Histogram of Pixel Intensities')\n",
    "axes[2].imshow(mean_div, cmap='turbo')\n",
    "axes[2].set_title('Mean of downsized images')\n",
    "axes[2].set_axis_off()\n",
    "axes[3].plot(hist_centers2, hist2, lw=2)\n",
    "axes[3].set_title('Histogram of Pixel Intensities')\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a8c3bc79",
   "metadata": {},
   "source": [
    "This cell defines the function to to calculate a compressed representation of the image data and the next cells runs the function and allow for visualization as before.  Usually no changes need to be made here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66f3e764",
   "metadata": {},
   "outputs": [],
   "source": [
    "def _dct2d(mtrx: np.array):\n",
    "    \"\"\"\n",
    "    Calculates 2D discrete cosine transform.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    mtrx\n",
    "        Input matrix.  \n",
    "        \n",
    "    Returns\n",
    "    -------    \n",
    "    Discrete cosine transform of the input matrix.\n",
    "    \"\"\"\n",
    "     \n",
    "    # Check if input object is 2D.\n",
    "    if mtrx.ndim != 2:\n",
    "        raise ValueError(\"Passed object should be a matrix or a numpy array with dimension of two.\")\n",
    "\n",
    "    return dct(dct(mtrx.T, norm='ortho').T, norm='ortho')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de1b28ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "Dis_Cos_Trans_mean = _dct2d(mean_div.T)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6493ebf",
   "metadata": {},
   "outputs": [],
   "source": [
    "hist, hist_centers = exposure.histogram(mean_div)\n",
    "hist2, hist_centers2 = exposure.histogram(Dis_Cos_Trans_mean)\n",
    "fig, axes = plt.subplots(1, 4, figsize=(12, 3))\n",
    "axes[0].imshow(mean_div, cmap='turbo')\n",
    "axes[0].set_title('Mean Image')\n",
    "axes[0].set_axis_off()\n",
    "axes[1].plot(hist_centers, hist, lw=2)\n",
    "axes[1].set_title('Histogram of Pixel Intensities')\n",
    "axes[2].imshow(Dis_Cos_Trans_mean, cmap='turbo', vmin=-0.1, vmax=0.3)\n",
    "axes[2].set_title('Discrete Cosine \\nTransform of Mean Image')\n",
    "axes[2].set_axis_off()\n",
    "axes[3].plot(hist_centers2, hist2, lw=2)\n",
    "axes[3].set_title('Histogram of Pixel Intensities')\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "feb03743",
   "metadata": {},
   "source": [
    "### 3.3 Functions\n",
    "This cell defines functions to inverse the dct function, estimate the illumination correction, and resize the images to their original shape.  Usually no changes need to be made here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a24ca9c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def _idct2d(mtrx: np.array):\n",
    "    \"\"\"\n",
    "    Calculates 2D inverse discrete cosine transform.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    mtrx\n",
    "        Input matrix.  \n",
    "        \n",
    "    Returns\n",
    "    -------    \n",
    "    Inverse of discrete cosine transform of the input matrix.\n",
    "    \"\"\"\n",
    "     \n",
    "    # Check if input object is 2D.\n",
    "    if mtrx.ndim != 2:\n",
    "        raise ValueError(\"Passed object should be a matrix or a numpy array with dimension of two.\")\n",
    " \n",
    "    return idct(idct(mtrx.T, norm='ortho').T, norm='ortho')\n",
    "\n",
    "def _shrinkageOperator(matrix, epsilon):\n",
    "    temp1 = matrix - epsilon\n",
    "    temp1[temp1 < 0] = 0\n",
    "    temp2 = matrix + epsilon\n",
    "    temp2[temp2 > 0] = 0\n",
    "    res = temp1 + temp2\n",
    "    return res\n",
    "\n",
    "def _inexact_alm_rspca_l1(\n",
    "    images, \n",
    "    lambda_flatfield, \n",
    "    if_darkfield, \n",
    "    lambda_darkfield, \n",
    "    optimization_tolerance, \n",
    "    max_iterations,\n",
    "    weight=None, \n",
    "    ):\n",
    "\n",
    "    if weight is not None and weight.size != images.size:\n",
    "            raise IOError('weight matrix has different size than input sequence')\n",
    "\n",
    "    # Shape variables\n",
    "    p = images.shape[0]           # Image height\n",
    "    q = images.shape[1]           # Image width\n",
    "    m = p*q                       # Total pixels per image\n",
    "    n = images.shape[2]          # Number of images\n",
    "    images = np.reshape(images, (m, n), order='F')\n",
    "\n",
    "    if weight is not None:\n",
    "        weight = np.reshape(weight, (m, n), order='F')\n",
    "    else:\n",
    "        weight = np.ones_like(images)\n",
    "\n",
    "    # SVD and Norm variables    \n",
    "    _, svd, _ = np.linalg.svd(images, full_matrices=False) # SVD decomposition\n",
    "    norm_two = svd[0] # Largest singular value\n",
    "    d_norm = np.linalg.norm(images, ord='fro') # Frobenius norm of images\n",
    "\n",
    "    # Optimization parameters\n",
    "    dual_var_lowrank = 0 # Dual variable for low-rank component (Y1)\n",
    "    lagrange_mult1 = 1 # Lagrange multiplier for first constraint (ent1)\n",
    "    lagrange_mult2 = 10 # Lagrange multiplier for second constraint (ent2)\n",
    "    penalty_factor = 12.5 / norm_two # (mu)\n",
    "    penalty_factor_bar = penalty_factor * 1e7 # Upper bound for penalty (mu_bar)\n",
    "    scale_ratio = 1.5 # Scale factor for penalty updates\n",
    "\n",
    "    # Component matrices\n",
    "    A1_hat = np.zeros_like(images) # Estimated flat-field\n",
    "    A1_coeff = np.ones((1, images.shape[1])) # Flat-field coefficients\n",
    "    E1_hat = np.zeros_like(images) # Estimated error/noise\n",
    "    W_hat = _dct2d(np.zeros((p, q)).T) # DCT coefficients\n",
    "\n",
    "    # Offset and mask variables\n",
    "    A_offset = np.zeros((m, 1)) # Offset for flatfield\n",
    "    B1_uplimit = np.min(images) # Upper limit for darkfield\n",
    "    B1_offset = 0 # Darkfield offset\n",
    "    A_inmask = np.zeros((p, q)) # Inner mask for optimization\n",
    "    # Mask covers central region (5/6 x 5/6)\n",
    "    A_inmask[int(np.round(p / 6) - 1): int(np.round(p*5 / 6)), int(np.round(q / 6) - 1): int(np.round(q * 5 / 6))] = 1\n",
    "\n",
    "    # main iteration loop starts\n",
    "    iter = 0\n",
    "    total_svd = 0\n",
    "    converged = False\n",
    "\n",
    "    while not converged:\n",
    " \n",
    "        iter += 1\n",
    "\n",
    "        if len(A1_coeff.shape) == 1:\n",
    "            A1_coeff = np.expand_dims(A1_coeff, 0)\n",
    "        if len(A_offset.shape) == 1:\n",
    "            A_offset = np.expand_dims(A_offset, 1)\n",
    "\n",
    "        W_idct_hat = _idct2d(W_hat.T)\n",
    "        A1_hat = np.dot(np.reshape(W_idct_hat, (-1,1), order='F'), A1_coeff) + A_offset\n",
    "\n",
    "        temp_W = (images - A1_hat - E1_hat + (1 / penalty_factor) * dual_var_lowrank) / lagrange_mult1\n",
    "        temp_W = np.reshape(temp_W, (p, q, n), order='F')\n",
    "        temp_W = np.mean(temp_W, axis=2)\n",
    "\n",
    "        W_hat = W_hat + _dct2d(temp_W.T)\n",
    "        W_hat = np.maximum(W_hat - lambda_flatfield / (lagrange_mult1 * penalty_factor), 0) + np.minimum(W_hat + lambda_flatfield / (lagrange_mult1 * penalty_factor), 0)\n",
    "        W_idct_hat = _idct2d(W_hat.T)\n",
    "\n",
    "        if len(A1_coeff.shape) == 1:\n",
    "            A1_coeff = np.expand_dims(A1_coeff, 0)\n",
    "        if len(A_offset.shape) == 1:\n",
    "            A_offset = np.expand_dims(A_offset, 1)\n",
    "\n",
    "        A1_hat = np.dot(np.reshape(W_idct_hat, (-1,1), order='F'), A1_coeff) + A_offset\n",
    "        E1_hat = images - A1_hat + (1 / penalty_factor) * dual_var_lowrank / lagrange_mult1\n",
    "        E1_hat = _shrinkageOperator(E1_hat, weight / (lagrange_mult1 * penalty_factor))\n",
    "        R1 = images - E1_hat\n",
    "        A1_coeff = np.mean(R1, 0) / np.mean(R1)\n",
    "        A1_coeff[A1_coeff < 0] = 0\n",
    "\n",
    "        if if_darkfield:\n",
    "            validA1coeff_idx = np.where(A1_coeff < 1)\n",
    "\n",
    "            B1_coeff = (np.mean(R1[np.reshape(W_idct_hat, -1, order='F') > np.mean(W_idct_hat) - 1e-6][:, validA1coeff_idx[0]], 0) - \\\n",
    "            np.mean(R1[np.reshape(W_idct_hat, -1, order='F') < np.mean(W_idct_hat) + 1e-6][:, validA1coeff_idx[0]], 0)) / np.mean(R1)\n",
    "            k = np.array(validA1coeff_idx).shape[1]\n",
    "            temp1 = np.sum(A1_coeff[validA1coeff_idx[0]]**2)\n",
    "            temp2 = np.sum(A1_coeff[validA1coeff_idx[0]])\n",
    "            temp3 = np.sum(B1_coeff)\n",
    "            temp4 = np.sum(A1_coeff[validA1coeff_idx[0]] * B1_coeff)\n",
    "            temp5 = temp2 * temp3 - temp4 * k\n",
    "            if temp5 == 0:\n",
    "                B1_offset = 0\n",
    "            else:\n",
    "                B1_offset = (temp1 * temp3 - temp2 * temp4) / temp5\n",
    "\n",
    "            B1_offset = np.maximum(B1_offset, 0)\n",
    "            B1_offset = np.minimum(B1_offset, B1_uplimit / np.mean(W_idct_hat))\n",
    "            B_offset = B1_offset * np.reshape(W_idct_hat, -1, order='F') * (-1)\n",
    "            B_offset = B_offset + np.ones_like(B_offset) * B1_offset * np.mean(W_idct_hat)\n",
    "\n",
    "            A1_offset = np.mean(R1[:, validA1coeff_idx[0]], axis=1) - np.mean(A1_coeff[validA1coeff_idx[0]]) * np.reshape(W_idct_hat, -1, order='F')\n",
    "            A1_offset = A1_offset - np.mean(A1_offset)\n",
    "            A_offset = A1_offset - np.mean(A1_offset) - B_offset\n",
    "\n",
    "            # smooth A_offset\n",
    "            W_offset = _dct2d(np.reshape(A_offset, (p,q), order='F').T)\n",
    "            W_offset = np.maximum(W_offset - lambda_darkfield / (lagrange_mult2 * penalty_factor), 0) + \\\n",
    "                np.minimum(W_offset + lambda_darkfield / (lagrange_mult2 * penalty_factor), 0)\n",
    "            A_offset = _idct2d(W_offset.T)\n",
    "            A_offset = np.reshape(A_offset, -1, order='F')\n",
    "\n",
    "            # encourage sparse A_offset\n",
    "            A_offset = np.maximum(A_offset - lambda_darkfield / (lagrange_mult2 * penalty_factor), 0) + \\\n",
    "                np.minimum(A_offset + lambda_darkfield / (lagrange_mult2 * penalty_factor), 0)\n",
    "            A_offset = A_offset + B_offset\n",
    "\n",
    "\n",
    "        Z1 = images - A1_hat - E1_hat # Constraint violation\n",
    "        dual_var_lowrank = dual_var_lowrank + penalty_factor * Z1 # Dual variable updates\n",
    "        penalty_factor = np.minimum(penalty_factor * scale_ratio, penalty_factor_bar)\n",
    "        \n",
    "        # Stop Criterion\n",
    "        stopCriterion = np.linalg.norm(Z1, ord='fro') / d_norm\n",
    "        print(f'Iteration {iter}, stopCriterion: {stopCriterion}')\n",
    "        if stopCriterion < optimization_tolerance:\n",
    "            converged = True\n",
    "\n",
    "        if not converged and iter >= max_iterations:\n",
    "            converged = True\n",
    "    A_offset = np.squeeze(A_offset)\n",
    "    A_offset = A_offset + B1_offset * np.reshape(W_idct_hat, -1, order='F')\n",
    "    return A1_hat, E1_hat, A_offset\n",
    "\n",
    "def _resize_image(image: np.ndarray, side_size: float  = None, x_side_size: float = None, y_side_size: float = None):\n",
    "    if side_size is not None:\n",
    "        y_side_size = x_side_size = side_size\n",
    "    if image.shape[0] != x_side_size or image.shape[1] != y_side_size:\n",
    "        return skresize(\n",
    "            image,\n",
    "            (x_side_size, y_side_size), \n",
    "            order = RESIZE_ORDER, \n",
    "            mode = RESIZE_MODE,\n",
    "            preserve_range = PRESERVE_RANGE\n",
    "        )\n",
    "    else:\n",
    "        return image"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46d7ff26",
   "metadata": {},
   "source": [
    "### 3.4 Estimate Correction\n",
    "The following cell runs the illumination correction estimation and returns two models.  It is based on the Basic correction algorithm implemented by the Peng lab https://github.com/peng-lab as a Cellprofiler plugin.  Here is where you may make changes to tune the algorithm for optimal results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2988abae",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Tunable parameters\n",
    "if_darkfield = True\n",
    "optimization_tolerance = 1e-6\n",
    "max_iterations = 500\n",
    "max_reweight_iterations = 25\n",
    "reweight_tolerance = 1.0e-3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0041e28a",
   "metadata": {},
   "outputs": [],
   "source": [
    "XAoffset = np.zeros((nrows, ncols))\n",
    "weight = np.ones(downsized_image.shape)\n",
    "flag_reweighting = True\n",
    "flatfield_last = np.ones((nrows, ncols))\n",
    "darkfield_last = np.random.randn(nrows, ncols)\n",
    "sorted_images = np.sort(downsized_image, axis=2)\n",
    "epsilon = 0.1\n",
    "lambda_flatfield = np.sum(np.abs(Dis_Cos_Trans_mean)) / 400 * 0.5\n",
    "lambda_darkfield = lambda_flatfield * 0.2\n",
    "reweighting_iter = 0\n",
    "\n",
    "while flag_reweighting:\n",
    "    reweighting_iter += 1\n",
    "\n",
    "    initial_flatfield = False\n",
    "    if initial_flatfield:\n",
    "        raise IOError('Initial flatfield option not implemented yet!')\n",
    "    else:\n",
    "        X_k_A, X_k_E, X_k_Aoffset = _inexact_alm_rspca_l1(\n",
    "            images = sorted_images, \n",
    "            lambda_flatfield = lambda_flatfield,\n",
    "            if_darkfield = if_darkfield, \n",
    "            lambda_darkfield = lambda_darkfield, \n",
    "            optimization_tolerance = optimization_tolerance, \n",
    "            max_iterations = max_iterations,\n",
    "            weight=weight\n",
    "        )\n",
    "\n",
    "    XA = np.reshape(X_k_A, [nrows, ncols, -1], order='F')\n",
    "    XE = np.reshape(X_k_E, [nrows, ncols, -1], order='F')\n",
    "    XAoffset = np.reshape(X_k_Aoffset, [nrows, ncols], order='F')\n",
    "    XE_norm = XE / np.mean(XA, axis=(0, 1))\n",
    "\n",
    "    weight = np.ones_like(XE_norm) / (np.abs(XE_norm) + epsilon)\n",
    "\n",
    "    weight = weight * weight.size / np.sum(weight)\n",
    "\n",
    "    temp = np.mean(XA, axis=2) - XAoffset\n",
    "    flatfield_current = temp / np.mean(temp)\n",
    "    darkfield_current = XAoffset\n",
    "    mad_flatfield = np.sum(np.abs(flatfield_current - flatfield_last)) / np.sum(np.abs(flatfield_last))\n",
    "    temp_diff = np.sum(np.abs(darkfield_current - darkfield_last))\n",
    "    \n",
    "    if temp_diff < 1e-7:\n",
    "        mad_darkfield = 0\n",
    "    else:\n",
    "        mad_darkfield = temp_diff / np.maximum(np.sum(np.abs(darkfield_last)), 1e-6)\n",
    "\n",
    "    print(f\"Re-weighting iteration {reweighting_iter}: MAD flatfield = {mad_flatfield}; MAD darkfield = {mad_darkfield}\")\n",
    "    flatfield_last = flatfield_current\n",
    "    darkfield_last = darkfield_current\n",
    "    if np.maximum(mad_flatfield,\n",
    "                    mad_darkfield) <= reweight_tolerance or \\\n",
    "            reweighting_iter >= max_reweight_iterations:\n",
    "        flag_reweighting = False\n",
    "\n",
    "shading = np.mean(XA, 2) - XAoffset\n",
    "flatfield = _resize_image(\n",
    "    image = shading, \n",
    "    x_side_size = _saved_size[0], \n",
    "    y_side_size = _saved_size[1]\n",
    ")\n",
    "flatfield = flatfield.astype(np.float64)  # Ensure floating point\n",
    "flatfield = np.clip(flatfield, 1e-6, None)  # Prevent division by zero\n",
    "flatfield = flatfield / np.mean(flatfield[flatfield > 0])\n",
    "\n",
    "if if_darkfield:\n",
    "    darkfield = _resize_image(\n",
    "        image = XAoffset, \n",
    "        x_side_size = _saved_size[0], \n",
    "        y_side_size = _saved_size[1]\n",
    "    )\n",
    "else:\n",
    "    darkfield = np.zeros_like(flatfield)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d1b26510",
   "metadata": {},
   "source": [
    "Visualize intermediate results as above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5cfee25f",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Pre re-weighting flatfield image shape is {XA.shape} and type is {XA.dtype}\\nMin value is {XA.min()} and max value is {XA.max()}\")\n",
    "print(f\"Error image shape is {XE.shape} and type is {XE.dtype}\\nMin value is {XE.min()} and max value is {XE.max()}\")\n",
    "hist, hist_centers = exposure.histogram(XA)\n",
    "hist2, hist_centers2 = exposure.histogram(XE)\n",
    "hist3, hist_centers3 = exposure.histogram(XAoffset)\n",
    "fig, axes = plt.subplots(1, 6, figsize=(16, 3))\n",
    "axes[0].imshow(XA[:,:,i], cmap='turbo')\n",
    "axes[0].set_title('Pre re-weighting flatfield') \n",
    "axes[0].set_axis_on()\n",
    "axes[1].plot(hist_centers, hist, lw=2)\n",
    "axes[1].set_title('Histogram of Pixel Intensities')\n",
    "axes[2].imshow(XE[:,:,i], cmap='turbo')\n",
    "axes[2].set_title('Optimized error')\n",
    "axes[2].set_axis_on()\n",
    "axes[3].plot(hist_centers2, hist2, lw=2)\n",
    "axes[3].set_title('Histogram of Pixel Intensities')\n",
    "axes[4].imshow(XAoffset, cmap='turbo')\n",
    "axes[4].set_title('Pre re-weighting darkfield')\n",
    "axes[4].set_axis_on()\n",
    "axes[5].plot(hist_centers3, hist3, lw=2)\n",
    "axes[5].set_title('Histogram of Pixel Intensities')\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ffbc20d4",
   "metadata": {},
   "source": [
    "Visualize the models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "077377df",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Flatfield image shape is {flatfield.shape} and type is {flatfield.dtype}\\nMin value is {flatfield.min()} and max value is {flatfield.max()}\")\n",
    "print(f\"Darkfield image shape is {darkfield.shape} and type is {darkfield.dtype}\\nMin value is {darkfield.min()} and max value is {darkfield.max()}\")\n",
    "hist, hist_centers = exposure.histogram(flatfield)\n",
    "hist2, hist_centers2 = exposure.histogram(darkfield)\n",
    "fig, axes = plt.subplots(1, 4, figsize=(12, 3))\n",
    "axes[0].imshow(flatfield, cmap='turbo')\n",
    "axes[0].set_title('Flatfield')\n",
    "axes[0].set_axis_on()\n",
    "axes[1].plot(hist_centers, hist, lw=2)\n",
    "axes[1].set_title('Histogram of Pixel Intensities')\n",
    "axes[2].imshow(darkfield, cmap='turbo')\n",
    "axes[2].set_title('Darkfield')\n",
    "axes[2].set_axis_on()\n",
    "axes[3].plot(hist_centers2, hist2, lw=2)\n",
    "axes[3].set_title('Histogram of Pixel Intensities')\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "750ec56d",
   "metadata": {},
   "source": [
    "### 3.5 Apply to Data\n",
    "If you are satisfied with the results, apply the correction to a copy of the raw data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce938188",
   "metadata": {},
   "outputs": [],
   "source": [
    "corrected = np.zeros_like(im_array_init, dtype=np.float64)\n",
    "\n",
    "for i in range(len(im_array_init)):\n",
    "    corrected[i] = (((im_array_init[i].astype(np.float64)) - darkfield) / flatfield)\n",
    "    corrected[i] = np.clip(corrected[i], im_array_init[i].min(), (im_array_init[i].max()))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e4367f92",
   "metadata": {},
   "source": [
    "Visualize the original and compare to the corrected image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05280941",
   "metadata": {},
   "outputs": [],
   "source": [
    "i=30\n",
    "print(f\"Original image shape is {im_array_init.shape} and type is {im_array_init.dtype}\\nMin value is {im_array_init.min()} and max value is {im_array_init.max()}\")\n",
    "print(f\"Corrected image shape is {corrected.shape} and type is {corrected.dtype}\\nMin value is {corrected.min()} and max value is {corrected.max()}\")\n",
    "hist, hist_centers = exposure.histogram(im_array_init[i])\n",
    "hist2, hist_centers2 = exposure.histogram(corrected[i])\n",
    "intensity_profile_init = np.mean(im_array_init[i], axis=0)\n",
    "intensity_profile_cor = np.mean(corrected[i], axis=0)\n",
    "dif = intensity_profile_init-intensity_profile_cor\n",
    "fig, axes = plt.subplots(1, 3, figsize=(12, 3))\n",
    "axes[0].imshow(im_array_init[i], cmap='turbo')\n",
    "axes[0].set_title('Original image')\n",
    "axes[0].set_axis_on()\n",
    "axes[1].imshow(corrected[i], cmap='turbo')\n",
    "axes[1].set_title('Corrected image')\n",
    "axes[1].set_axis_on()\n",
    "axes[2].plot(dif)\n",
    "axes[2].set_title('Difference in Mean Pixel Intensities')\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "231c895c-bdee-4e4e-aae3-061b0903b591",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "## 4. Stitching"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae05c0bb-751e-404c-a6d1-6a57a58fcabe",
   "metadata": {},
   "source": [
    "### 4.1 Find Stitching Positions and Save Model\n",
    "- Here we stitch the original and corrected tiles for comparison.  This is accomplished using the package m2stitch https://github.com/yfukai/m2stitch.  \n",
    "- Enter the number of rows and columns, and overlap percentage known from your data.   \n",
    "- The rows and columns functions are for generating lists corresponding to a \"snake by rows\" pattern starting at the upper left of the image.  They will need to be changed for other stitching patterns.\n",
    "- Enter use_gpu=True to use your gpu; False to use CPU.\n",
    "- To save and reuse stitching results, uncomment/comment out the appropriate lines.  The results are saved to the metadata folder defined above.\n",
    "- Be sure to closely evaluate the areas of overlap for stitching quality.  The cells below allow for visualizing in this notebook and for saving to inspect elsewhere.\n",
    "- Adjust pou (percent overlap uncertainty) and/or overlap_percentage to affect stitching quality. \n",
    "- For more information see https://github.com/yfukai/m2stitch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1261aed6-bde7-4b27-88d9-653a11377ece",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "n = 9  # Number of rows (height)\n",
    "m = 7 # Number of columns (width)\n",
    "overlap_percentage = 0.30\n",
    "\n",
    "pou = 6\n",
    "stitch_model = os.path.join(meta_dir, \"result.pkl\")\n",
    "\n",
    "# Row coordinates: each row index is repeated m times\n",
    "rows = list(chain.from_iterable(repeat(row, m) for row in range(n)))\n",
    "\n",
    "# Column coordinates: snake pattern for each row, going back and forth\n",
    "cols = list(chain.from_iterable(\n",
    "    range(m) if row % 2 == 0 else range(m - 1, -1, -1) for row in range(n)\n",
    "))\n",
    "result_df, _ = stitch_images(im_array_init, rows, cols, initial_ncc_threshold = 0, overlap_percentage=overlap_percentage, pou=pou, max_cores=10, use_gpu=True)\n",
    "result_df.to_pickle(stitch_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "796ee940-8ef9-4ca9-90f2-c0400defdd26",
   "metadata": {},
   "source": [
    "### Section 4.2\n",
    "If the cell above has already been run to create a stitching model, run the cell below to apply it to the original images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4df33a5-9791-4448-85c4-f28d4e3ad7b1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "stitch_model = os.path.join(meta_dir, \"result.pkl\")\n",
    "\n",
    "result_df = pd.read_pickle(stitch_model)\n",
    "\n",
    "result_df[\"y_pos2\"] = result_df[\"y_pos\"] - result_df[\"y_pos\"].min()\n",
    "result_df[\"x_pos2\"] = result_df[\"x_pos\"] - result_df[\"x_pos\"].min()\n",
    "\n",
    "size_y = im_array_init.shape[1]\n",
    "size_x = im_array_init.shape[2]\n",
    "\n",
    "stitched_image_size = (\n",
    "    result_df[\"y_pos2\"].max() + size_y,\n",
    "    result_df[\"x_pos2\"].max() + size_x,\n",
    ")\n",
    "stitched_image = np.zeros_like(im_array_init, shape=stitched_image_size)\n",
    "for i, row in result_df.iterrows():\n",
    "    stitched_image[\n",
    "        row[\"y_pos2\"] : row[\"y_pos2\"] + size_y,\n",
    "        row[\"x_pos2\"] : row[\"x_pos2\"] + size_x,\n",
    "    ] = im_array_init[i]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db018213-9ff0-4657-9ec3-02c6e4677f7b",
   "metadata": {},
   "source": [
    "- Visualize the result of stitching the original tiles.  \n",
    "- There should be a toolbar to the left of the image when hovering over it.  If there isn't, be sure the appropriate renderer is chosen by clicking the ellipses next to the plot and choose “Change Presentation.” It should be “Jupyter IPyWidget Renderer.”  This will allow you to zoom-to-rectangle, resize, save, and move between views."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8228d715-e449-4611-8ab7-32ca159a5462",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "Kview.imshow(stitched_image, colormap='turbo')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "abc0990e-79ca-4e54-a2a3-42862d01e6cf",
   "metadata": {},
   "source": [
    "Save the stitched image if desired."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb846537-be67-4824-a7d7-3d3844b52212",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "raw_image_file_path = os.path.join(meta_dir, \"test_original.tif\") \n",
    "imsave(raw_image_file_path, stitched_image)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73ad4d25-24c2-4c67-b609-3e144d681bd4",
   "metadata": {},
   "source": [
    "### Section 4.3\n",
    "Reuse the stitching model to stitch the corrected tiles."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6984ccb8-c7a8-4511-8af8-0a917f922253",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "result_df=pd.read_pickle(stitch_model)\n",
    "result_df[\"y_pos2\"] = result_df[\"y_pos\"] - result_df[\"y_pos\"].min()\n",
    "result_df[\"x_pos2\"] = result_df[\"x_pos\"] - result_df[\"x_pos\"].min()\n",
    "\n",
    "size_y = corrected.shape[1]\n",
    "size_x = corrected.shape[2]\n",
    "\n",
    "stitched_image_size = (\n",
    "    result_df[\"y_pos2\"].max() + size_y,\n",
    "    result_df[\"x_pos2\"].max() + size_x,\n",
    ")\n",
    "stitched_image_corrected = np.zeros_like(corrected, shape=stitched_image_size)\n",
    "for i, row in result_df.iterrows():\n",
    "    stitched_image_corrected[\n",
    "        row[\"y_pos2\"] : row[\"y_pos2\"] + size_y,\n",
    "        row[\"x_pos2\"] : row[\"x_pos2\"] + size_x,\n",
    "    ] = corrected[i]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af7325c3-1562-42b8-9b03-7647c8f86061",
   "metadata": {},
   "source": [
    "Visualize the result of stitching the corrected tiles. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42d64ca7-9e63-44dc-9f33-6b032a74ae64",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "plt.figure() \n",
    "stackview.imshow(stitched_image_corrected, colormap='turbo')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d052bc5-3940-4aa9-be44-e32a92ca4fbd",
   "metadata": {},
   "source": [
    "Save the corrected stitched image if desired."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e79273c7-bd3f-4e74-8fa8-984a34f5020f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "result_image_file_path = os.path.join(meta_dir, \"test_corrected.tif\") \n",
    "imsave(result_image_file_path, stitched_image_corrected, check_contrast=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a6c20b6c-76de-4d16-8b47-d9ce92100233",
   "metadata": {},
   "source": [
    "Before testing deconvolution, you must correct and stitch an entire z-stack.  For this you can run the cells above for each image in the stack, or use the notebook for batch processing: 2_Cycle_Processing.ipynb"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee52a231-4fc2-455d-abfa-0f9e97c6eb3a",
   "metadata": {},
   "source": [
    "## 5. Deconvolution"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e143f9c-2371-46bb-abc4-eed1fe016590",
   "metadata": {},
   "source": [
    "- The original deconvolution program was written in Matlab for lightsheet images.  It has been extensively modified for use with widefield fluorescence.  See https://www.nature.com/articles/s41598-019-53875-y \n",
    "\n",
    "- Enter the necessary parameters in the cell below.  The \"iterations\" and \"stop_criterion\" values can be tuned to get good results.  The stop criteria will be the % mean squared error between iterations. The damping and hist_clip further can improve results by removing noise and limiting extreme pixel values, respectively.  \n",
    "\n",
    "- There is no need to supply a point-spread function.  \n",
    "\n",
    "- Finding the optimal use of computing resources on GPU or CPU is accomplished via the max_memory and blocksize parameters.  Max memory will be defined based on 80% of physical memory for GPU.  You may specify a blocksize that fits with your device architecture and dataset (multiples of 32 work well) or set to \"1\" for an automatic determination.  Automatic determination requires a lower (i.e. < 0.50) max_memory. Using these methods, there is no practical limit to image size.\n",
    "\n",
    "- Be sure the MATLAB path is properly defined.\n",
    "\n",
    "- The image_dir and stitch_dir definitions used previously will be again used here, and a new folder will be created to save the deconvolved images.  \n",
    "\n",
    "- If there is an \"Maximum variable size allowed on the device is exceeded.\" error, then restart the kernel, decrease the max_memory parameter, and try again."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "beb6c3b5-66b0-4fb9-bfc3-6e10c21b2f8e",
   "metadata": {},
   "outputs": [],
   "source": [
    "dec_cycles = 1\n",
    "dec_channels = 1\n",
    "\n",
    "# pixel size in xy dimension (nanometers)\n",
    "xy_vox = 377\n",
    "# pixel size in z dimension (nanometers)\n",
    "z_vox = 1500\n",
    "# Number of iterations of Lucy-Richardson algo before stopping unless stop_crit is met first\n",
    "iterations = 25\n",
    "# Microscope objective numerical aperture\n",
    "mic_NA = 0.75\n",
    "# Refractive index of tissue being imaged\n",
    "tissue_RI = 1.5\n",
    "# Opening size in millimeters of objective aperture\n",
    "slit_aper = 6.5\n",
    "# Focal length in millimeters of objective\n",
    "f_cyl = 1\n",
    "# Used to reduce noise.  Increase value for noisy images. (0-10)\n",
    "damping = 0\n",
    "# If set, the deconvolved images will be clipped by this percent for max and min values, and then scaled to full range of bit depth. (0-5)\n",
    "hist_clip = 0.010\n",
    "# Percent change between iterations to use as criteria to stop deconvolution.\n",
    "stop_criterion = 5.00\n",
    "# Percent maximum GPU or CPU memory \n",
    "max_memory = 0.5\n",
    "# Enter 'GPU' or 'CPU'\n",
    "device = 'GPU'\n",
    "# Chunk size in bytes\n",
    "blocksize = 8196\n",
    "# The excitation and emission wavelength in nanometers\n",
    "ex = 358\n",
    "em = 461\n",
    "\n",
    "decon_dir = stitch_dir.replace('_BaSiC_Stitched', '_Decon')\n",
    "source = os.path.join(stitch_dir, f\"cyc{str(dec_cycles).zfill(2)}\", f\"CH{str(dec_channels)}\")\n",
    "dest = os.path.join(decon_dir, f\"cyc{str(dec_cycles).zfill(2)}\", f\"CH{str(dec_channels)}\")\n",
    "os.makedirs(dest, exist_ok=True)\n",
    "\n",
    "command = [source, str(xy_vox), str(z_vox), str(iterations), str(mic_NA), str(tissue_RI), str(ex), str(em), str(f_cyl), str(slit_aper), str(damping), str(hist_clip), str(stop_criterion), str(max_memory), device, str(blocksize)]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8c056e1",
   "metadata": {},
   "source": [
    "This cell defines the function to enable output of the Matlab script as it is processing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "7e1f9381",
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_subprocess(command):\n",
    " \n",
    "    process = subprocess.Popen(command, \n",
    "                                stdout=subprocess.PIPE, \n",
    "                                stderr=subprocess.STDOUT, \n",
    "                                shell=True, \n",
    "                                bufsize=1,\n",
    "                                universal_newlines=True, \n",
    "                                text=True)\n",
    "    \n",
    "    while True:\n",
    "        output = process.stdout.readline()\n",
    "        if output == '' and process.poll() is not None:\n",
    "            break\n",
    "        if output:\n",
    "            print(output, end='') \n",
    "\n",
    "    rc = process.poll()\n",
    "    return rc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e0d0c3b",
   "metadata": {},
   "outputs": [],
   "source": [
    "mathome = os.path.join(base_dir, \"KINTSUGI\", \"R2024a\", \"runtime\", \"win64\")\n",
    "os.environ['PATH'] = mathome + os.pathsep + os.environ['PATH']\n",
    "decon_exe = os.path.join(base_dir, \"KINTSUGI\", \"KDecon_v4.exe\") \n",
    "command.insert(0, decon_exe) \n",
    "\n",
    "run_subprocess(command) \n",
    "os.rename(os.path.join(source, 'deconvolved'), os.path.join(dest, 'deconvolved'))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "016a6de1-9bd3-47ad-8e56-d881d6828a26",
   "metadata": {},
   "source": [
    "The following cell will compare the original to the deconvolved image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2281a00-4955-44a1-b044-88b12986ced8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Enter the cycle, channel, and z-plane for the image you want to evaluate.\n",
    "c = 1\n",
    "ch = 1\n",
    "z = 8\n",
    "\n",
    "# Crop the image by entering the x and y coordinates below.\n",
    "x1 = 600\n",
    "x2 = 1000\n",
    "y1 = 800\n",
    "y2 = 1200\n",
    "\n",
    "source_image = imread(os.path.join(stitch_dir, f\"cyc{str(c).zfill(2)}\", f\"CH{str(ch)}\", f\"{str(z).zfill(2)}.tif\"))\n",
    "dest = os.path.join(stitch_dir.replace('_BaSiC_Stitched', '_Decon'), f\"cyc{str(c).zfill(2)}\", f\"CH{str(ch)}\")\n",
    "fig, axes = plt.subplots(1, 2, figsize=(12, 9))\n",
    "im = axes[0].imshow(source_image[x1:x2,y1:y2], cmap='magma')\n",
    "fig.colorbar(im, shrink=0.5, ax=axes[0])\n",
    "axes[0].set_title(\"Original\")\n",
    "decon_image = imread(os.path.join(dest, 'deconvolved', f\"deconv_{str(z).zfill(3)}.tif\"))\n",
    "im = axes[1].imshow(decon_image[x1:x2,y1:y2], cmap='magma')\n",
    "fig.colorbar(im, shrink=0.5, ax=axes[1])\n",
    "axes[1].set_title(\"Deconvolved\")\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "85fa4f09-712c-47d5-ab1d-c5cacf83edde",
   "metadata": {},
   "source": [
    "## 6. FIJI Clij2 plugin - EDF"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a373bd03-dfcc-444a-ad9a-4426180dc486",
   "metadata": {},
   "source": [
    "### 6.1 Define FIJI variables\n",
    "- The following section connects this notebook with a local installation of FIJI and a Clij2 plugin to run an Extended Depth of Focus variance projection macro along the deconvolved image stack.  \n",
    "- The image calculation is best done on the GPU for efficient (~10x faster) processing although any OpenCL device may be used.  \n",
    "- In the next cell, define path values and enter the amount of RAM to \"ij_mem\" (no more than 80% of system RAM is a good rule of thumb) to allow for the computation.  \n",
    "- The following cells will list possible devices to use; the name of the device chosen must match exactly for \"device.\"  \n",
    "- Note that if a CPU is used, an additional OpenCL driver may need to be installed (for example Intel Processors: https://www.intel.com/content/www/us/en/developer/articles/tool/opencl-drivers.html).  \n",
    "- A new folder will be created for the results, and the image can be inspected in another program and in the last cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf62c19c",
   "metadata": {},
   "outputs": [],
   "source": [
    "fiji_dir = os.path.join(base_dir, \"Fiji.app\")\n",
    "javahome = os.path.join(base_dir, \"KINTSUGI\", \"jdk-21_windows-x64_bin\", \"jdk-21.0.5\", \"bin\")\n",
    "os.environ['PATH'] = javahome + os.pathsep + os.environ['PATH']\n",
    "mavenhome = os.path.join(base_dir, 'KINTSUGI','apache-maven-3.9.9-bin','apache-maven-3.9.9', 'bin')\n",
    "os.environ['PATH'] = mavenhome + os.pathsep + os.environ['PATH']\n",
    "\n",
    "ij_mem = 50\n",
    "scyjava.config.add_option(f'-Xmx{str(ij_mem)}g')\n",
    "ij = imagej.init(fiji_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a044839a",
   "metadata": {},
   "source": [
    "### 6.2 Get Device Names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa2cc856",
   "metadata": {},
   "outputs": [],
   "source": [
    "platforms = cl.get_platforms()\n",
    "for platform in platforms:\n",
    "    print(f\"Platform: {platform.name}\")\n",
    "    devices = platform.get_devices()\n",
    "    for device in devices:\n",
    "        print(f\"  Device: {device.name} - Type: {cl.device_type.to_string(device.type)}\")\n",
    "\n",
    "# Select a CPU device (assuming the first CPU device found)\n",
    "cpu_device = None\n",
    "for platform in platforms:\n",
    "    devices = platform.get_devices()\n",
    "    for device in devices:\n",
    "        if device.type == cl.device_type.CPU:\n",
    "            cpu_device = device\n",
    "            break\n",
    "    if cpu_device:\n",
    "        break"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f6c39b4",
   "metadata": {},
   "source": [
    "### 6.3 Define Channel Names\n",
    "Enter channel names below.  The EDF output images will be named accordingly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7306698",
   "metadata": {},
   "outputs": [],
   "source": [
    "channel_name_dict = {\"cyc01.tif\" : [\"DAPI\", \"Blank1a\", \"Blank1b\", \"Blank1c\"],\n",
    " \"cyc02.tif\" : [\"DAPI\", \"CD31\", \"CD8\", \"Empty2c\"],\n",
    " \"cyc03.tif\" : [\"DAPI\", \"CD20\", \"Ki67\", \"CD3e\"],\n",
    " \"cyc04.tif\" : [\"DAPI\", \"SMActin\", \"Podoplanin\", \"CD68\"],\n",
    " \"cyc05.tif\" : [\"DAPI\", \"PanCK\", \"CD21\", \"CD4\"],\n",
    " \"cyc06.tif\" : [\"DAPI\", \"Lyve1\", \"CD45RO\", \"CD11c\"],\n",
    " \"cyc07.tif\" : [\"DAPI\", \"CD35\", \"ECAD\", \"CD107a\"],\n",
    " \"cyc08.tif\" : [\"DAPI\", \"CD34\", \"CD44\", \"HLADR\"],\n",
    " \"cyc09.tif\" : [\"DAPI\", \"Empty9a\", \"FoxP3\", \"CD163\"],\n",
    " \"cyc10.tif\" : [\"DAPI\", \"Empty10a\", \"CollagenIV\", \"Vimentin\"],\n",
    " \"cyc11.tif\" : [\"DAPI\", \"Empty11a\", \"CD15\", \"CD45\"],\n",
    " \"cyc12.tif\" : [\"DAPI\", \"Empty12a\", \"CD5\", \"CD1c\"],\n",
    " \"cyc13.tif\" : [\"DAPI\", \"Blank13a\", \"Blank13b\", \"Blank13c\"]}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f43f1a08",
   "metadata": {},
   "source": [
    "### 6.4 Run EDF\n",
    "Enter the cycle and channel for the image you want to evaluate."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9e90cda",
   "metadata": {},
   "outputs": [],
   "source": [
    "edf_cycles = 1\n",
    "edf_channels = 1\n",
    "\n",
    "decon_dir = stitch_dir.replace('_BaSiC_Stitched', '_Decon')\n",
    "edf_dir = decon_dir.replace('_Decon', '_EDF')\n",
    "edf_source = os.path.join(decon_dir, f\"cyc{str(edf_cycles).zfill(2)}\", f\"CH{str(edf_channels)}\", \"deconvolved\")\n",
    "edf_dest = os.path.join(edf_dir, f\"cyc{str(edf_cycles).zfill(2)}\")\n",
    "file_name = channel_name_dict.get(f\"cyc{str(edf_cycles).zfill(2)}.tif\")[edf_channels-1]\n",
    "os.makedirs(edf_dest, exist_ok=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "29fd8ae8",
   "metadata": {},
   "source": [
    "- Enter the OpenCl device name determined above.  If the device is not a GPU, enter \"False\" for GPU.\n",
    "- Enter the radii and sigma value for blurring.\n",
    "- Enter the number of tiles for splitting in the x and y dimensions (CPU will generally need many more e.g. if GPU split is 3 in each dimension, CPU splits will be 20).\n",
    "- If there are z-planes at the start or end of the stack that need to be removed or do not contribute to any in-focus pixels, they may be removed by setting z_start and z_end."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ddb7f713",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = \"NVIDIA RTX A4500\"\n",
    "\n",
    "radius_x = 5\n",
    "radius_y = 5\n",
    "sigma = 20.0\n",
    "\n",
    "xSplit = 3\n",
    "ySplit = 3\n",
    "\n",
    "z_start = 3\n",
    "z_end = 15\n",
    "\n",
    "GPU = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8e88d94-1767-4ad9-9125-d8cfe2dcf9e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "macro = \"\"\"\n",
    "#@ File in_folder\n",
    "#@ String device\n",
    "#@ File out_folder\n",
    "#@ String file_name\n",
    "#@ Integer radius_x\n",
    "#@ Integer radius_y\n",
    "#@ Float sigma\n",
    "#@ Integer xTiles\n",
    "#@ Integer yTiles\n",
    "#@ String z_start\n",
    "#@ String z_end\n",
    "#@ Boolean GPU\n",
    "\n",
    "run(\"CLIJ2 Macro Extensions\", \"cl_device=[\" + device + \"] automatic_output_naming=false\");\n",
    "Ext.CLIJ2_clear();\n",
    "//setBatchMode(true);\n",
    "\n",
    "File.openSequence(in_folder);\n",
    "Stack.getDimensions(width, height, channels, slices, frames);\n",
    "\n",
    "if (slices != z_end - z_start + 1) {\n",
    "    run(\"Slice Keeper\", \"first=\" + z_start + \" last=\" + z_end + \" increment=1\");\n",
    "    input = \"deconvolved kept stack\";\n",
    "    Stack.getDimensions(width, height, channels, slices, frames);\n",
    "    } else {\n",
    "    input = \"deconvolved\";\n",
    "}\n",
    "\n",
    "if (GPU == true) {\n",
    "\toutput = \"output\";\n",
    "\tnewImage(output, \"16-bit black\", width, height, 1);\n",
    "    } else {\n",
    "    Ext.CLIJ2_push(input);\n",
    "\tExt.CLIJ2_create2D(output, width, height, 16);\n",
    "}\n",
    "\n",
    "numTilesZ = 1;\n",
    "numTilesX = xTiles;\n",
    "numTilesY = yTiles;\n",
    "\n",
    "tileDepth = round((slices/numTilesZ));\n",
    "tileWidth = round((width/numTilesX));\n",
    "tileHeight = round((height/numTilesY));\n",
    "\n",
    "for (x = 0; x < numTilesX; x++) {\n",
    "\tfor (y = 0; y < numTilesY; y++) {\n",
    "\t\tfor (z = 0; z < numTilesZ; z++) {\n",
    "\n",
    "\t\t\tExt.CLIJ2_pushTile(input, x, y, z, tileWidth, tileHeight, tileDepth, 0, 0, 0);\t\t\n",
    "\t\t\tExt.CLIJ2_extendedDepthOfFocusVarianceProjection(input, output, radius_x, radius_y, sigma);\n",
    "            Ext.CLIJ2_release(input);\n",
    "            if (GPU == true) {\n",
    "\t\t\t\tExt.CLIJ2_pullTile(output, x, y, z, tileWidth, tileHeight, 1, 0, 0, 0);\n",
    "\t\t\t\tExt.CLIJ2_release(output);\n",
    "\t\t\t}            \n",
    "\t\t}\n",
    "\t}\n",
    "}\n",
    "\n",
    "if (GPU == false) {\n",
    "\tExt.CLIJ2_pull(output);\n",
    "    Ext.CLIJ2_release(output);\n",
    "}\n",
    "\n",
    "selectImage(output);\n",
    "saveAs(\"Tiff\", out_folder + File.separator + file_name);\n",
    "\n",
    "Ext.CLIJ2_clear();\n",
    "run(\"Close All\");\n",
    "\"\"\"\n",
    "args = {'in_folder': edf_source, \n",
    "\t\t\"device\" : device, \n",
    "\t\t\"out_folder\" : edf_dest, \n",
    "\t\t\"file_name\" : file_name, \n",
    "\t\t\"radius_x\" : radius_x, \n",
    "\t\t\"radius_y\" : radius_y, \n",
    "\t\t\"sigma\" : sigma, \n",
    "\t\t\"xTiles\" : xSplit, \n",
    "\t\t\"yTiles\" : ySplit, \n",
    "\t\t\"z_start\" : z_start, \n",
    "\t\t\"z_end\" : z_end, \n",
    "\t\t\"GPU\" : GPU}\n",
    "\n",
    "ij.py.run_macro(macro, args)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b627d4c8-90ff-4871-958a-63cd59b71981",
   "metadata": {},
   "source": [
    "Visualize the results with periodic z-planes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24b8933b-b171-4352-8ec4-e2301ca4510d",
   "metadata": {},
   "outputs": [],
   "source": [
    "decon_stack = sorted(glob(os.path.join(edf_source, f\"deconv_0??.tif\")), key=alphanumeric_key)\n",
    "decon_stack = io.imread_collection(decon_stack)\n",
    "decon_image = np.asarray(decon_stack)\n",
    "edf_result = imread(os.path.join(edf_dest, file_name + \".tif\"))\n",
    "\n",
    "x1 = 2400\n",
    "x2 = 3000\n",
    "y1 = 2400\n",
    "y2 = 3000\n",
    "\n",
    "fig, axes = plt.subplots(2, 3, figsize=(12, 8))\n",
    "im = axes[0,0].imshow(decon_image[6][y1:y2, x1:x2], cmap='turbo')\n",
    "axes[0,0].set_title(\"decon_z007\")\n",
    "im = axes[0,1].imshow(decon_image[7][y1:y2, x1:x2], cmap='turbo')\n",
    "axes[0,1].set_title(\"decon_z008\")\n",
    "im = axes[0,2].imshow(decon_image[8][y1:y2, x1:x2], cmap='turbo')\n",
    "axes[0,2].set_title(\"decon_z009\")\n",
    "im = axes[1,0].imshow(decon_image[9][y1:y2, x1:x2], cmap='turbo')\n",
    "axes[1,0].set_title(\"decon_z010\")\n",
    "im = axes[1,1].imshow(decon_image[10][y1:y2, x1:x2], cmap='turbo')\n",
    "axes[1,1].set_title(\"decon_z011\")\n",
    "im = axes[1,2].imshow(edf_result[y1:y2, x1:x2], cmap='turbo')\n",
    "axes[1,2].set_title(\"EDF_var_radius_x5_radius_y5_sigma20\")\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c01df92f-7e3f-4d05-a4ad-60f77346ff68",
   "metadata": {},
   "source": [
    "Visualize the results as a full image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd005e22-c399-4504-97cd-5e5b546b03b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "edf_result = imread(os.path.join(edf_dest, file_name + \".tif\"))\n",
    "plt.figure()\n",
    "Kview.imshow(edf_result, colormap='turbo')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72af5c0e",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "KINTSUGI",
   "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.20"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}