miPolScope_fig1_u2os_labelfree.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "3075465c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import tifffile\n",
    "import napari\n",
    "\n",
    "from scipy.io import loadmat\n",
    "from skimage import transform\n",
    "\n",
    "from waveorder import waveorder_microscopy, wavelet_softThreshold\n",
    "from waveorder.io import WaveorderReader"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e59b977",
   "metadata": {},
   "source": [
    "# Table of contents\n",
    "- View dataset\n",
    "- Load data\n",
    "    - Load raw images\n",
    "    - Load calibration data\n",
    "    - Load background images\n",
    "- Recostruct label-free channels\n",
    "    - Register images\n",
    "        - Crop edges\n",
    "        - View registered images\n",
    "    - Reconstruct Stokes images\n",
    "        - Initialize reconstructor\n",
    "        - Reconstruct Raw Stokes images\n",
    "            - View raw Stokes images\n",
    "        - Denoise S0 images\n",
    "            - Compare raw and denoised images\n",
    "        - Normalize Stokes images\n",
    "            - View normalized Stokes images and background images\n",
    "        - Correct background\n",
    "            - View background corrected Stokes images\n",
    "        - Denoise S1 and S2 images\n",
    "            - View denoised Stokes images\n",
    "    - Compute transmission, phase, retardance, and orientation\n",
    "- View reconstruction results"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1186c30a-8fd8-45c0-b83f-8a0d5c632694",
   "metadata": {
    "tags": []
   },
   "source": [
    "# View dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "021d904b-6ea4-487c-bd06-6847ff1f9b84",
   "metadata": {},
   "outputs": [],
   "source": [
    "# assume data is in ~/Downloads folder\n",
    "data_dir = os.path.join(os.path.expanduser('~'), 'Downloads')\n",
    "data_path = os.path.join(data_dir, 'miPolScope_fig1_u2os_labelfree.zarr')\n",
    "\n",
    "# Check that data path exists\n",
    "if not os.path.exists(data_path):\n",
    "    raise ValueError('Data path does not exist.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "6771ca36-e5d8-41ba-aa73-bfa28e068273",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "version mismatch: detected:FormatV01, requested:FormatV03\n",
      "version mismatch: detected:FormatV03, requested:FormatV01\n"
     ]
    }
   ],
   "source": [
    "viewer = napari.Viewer()\n",
    "layers = viewer.open(os.path.join(data_path, 'Row_0/Col_0/Pos_000'), plugin='napari-ome-zarr')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a4f5d6bd-5a1b-4cea-9570-f0f768564eac",
   "metadata": {},
   "source": [
    "# Load data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e004ffe-ae7d-4f77-885a-d64ee331aaac",
   "metadata": {},
   "source": [
    "## Load raw images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "1f47cf23",
   "metadata": {},
   "outputs": [],
   "source": [
    "wo_data = WaveorderReader(data_path, data_type='zarr')\n",
    "I = wo_data.get_zarr(0)\n",
    "n_timepoints, n_channels, n_slices, *img_size = I.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "9ed34c5c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of time points: 41\n",
      "Number of channels: 8\n",
      "Number of slices: 40\n",
      "Image size: [640, 640]\n"
     ]
    }
   ],
   "source": [
    "print(\n",
    "    f'Number of time points: {n_timepoints}\\n'\n",
    "    f'Number of channels: {n_channels}\\n'\n",
    "    f'Number of slices: {n_slices}\\n'\n",
    "    f'Image size: {img_size}'\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c8a1f0aa-e1c0-4c42-a4b5-5ce1e8304590",
   "metadata": {},
   "source": [
    "As demonstration, we will analyze only the first 5 timepoints"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "03496287-baef-4311-b3ae-b5209f8499e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "n_timepoints = 5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "5fb6a009-ada4-410b-912f-37af840655b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# load data into memory\n",
    "I = np.array(I[:n_timepoints, :4])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "898f10d2",
   "metadata": {},
   "source": [
    "## Load calibration data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "e1222df0",
   "metadata": {},
   "outputs": [],
   "source": [
    "cal_data = loadmat(os.path.join(data_path,'calibration.mat'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "cc40e956",
   "metadata": {},
   "outputs": [],
   "source": [
    "A = np.transpose(cal_data['A'].astype('float'), (2, 3, 0, 1)) # A has shape (size_Y, size_X, N_channels, N_Stokes)\n",
    "black_level = cal_data['black_level'][0][0].astype('uint16')\n",
    "\n",
    "tform0 = transform.AffineTransform(cal_data['tform0'].T)\n",
    "tform45 = transform.AffineTransform(cal_data['tform45'].T)\n",
    "tform90 = transform.AffineTransform(cal_data['tform90'].T)\n",
    "tform135 = transform.AffineTransform(cal_data['tform135'].T)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10b067da",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Load background images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "ff00c72e",
   "metadata": {},
   "outputs": [],
   "source": [
    "S_bg = tifffile.imread(os.path.join(data_path, 'Stokes_bg.ome.tif'))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5765e06",
   "metadata": {},
   "source": [
    "# Recostruct label-free channels"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12c9dbd8",
   "metadata": {},
   "source": [
    "## Register images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "19e47c3a",
   "metadata": {},
   "outputs": [],
   "source": [
    "I_registered = np.zeros((n_timepoints, 4, n_slices, *img_size), dtype='float')\n",
    "\n",
    "for t in range(n_timepoints):\n",
    "    for c, tform in enumerate((tform0, tform45, tform90, tform135)):\n",
    "        for z in range(n_slices):\n",
    "            I_registered[t,c,z] = transform.warp(I[t,c,z], tform.inverse, preserve_range=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b5853512",
   "metadata": {},
   "source": [
    "### Crop edges"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "3a8dcfb2",
   "metadata": {},
   "outputs": [],
   "source": [
    "I_registered = I_registered[..., 20:-20, 20:-20]\n",
    "img_size = I_registered.shape[-2:]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38a8dba2",
   "metadata": {},
   "source": [
    "### View registered images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "4a5015fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "viewer = napari.view_image(I_registered)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08b37a26",
   "metadata": {},
   "source": [
    "## Reconstruct Stokes images"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33368467-4922-495f-a70b-843b3d72d3c3",
   "metadata": {},
   "source": [
    "### Initialize reconstructor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "e52dbb5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "wavelength = 525 # in nm\n",
    "NA_obj = 1.2 # Numerical Aperture of Objective\n",
    "NA_illu = 0.4 # Numerical Aperture of Condenser\n",
    "n_objective_media = 1.33 # refractive index of objective immersion media\n",
    "mag = 30 # effective magnification\n",
    "n_slices = I.shape[-3] # number of slices in z-stack\n",
    "z_step_um = 0.25 # z-step size in um\n",
    "pad_z = 5 # slices to pad for phase reconstruction boundary artifacts\n",
    "pixel_size_um = 3.45 # camera pixel size in um\n",
    "bg_correction = 'local_fit' # BG correction method: 'None', 'local_fit', 'global'\n",
    "mode = '3D' # phase reconstruction mode, '2D' or '3D'\n",
    "use_gpu = False\n",
    "gpu_id = 0 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "34fb5853",
   "metadata": {},
   "outputs": [],
   "source": [
    "z_defocus = -(np.r_[:n_slices] - n_slices // 2) * z_step_um # assumes stack starts from the bottom\n",
    "swing = 0\n",
    "ps = pixel_size_um / mag\n",
    "\n",
    "reconstructor = waveorder_microscopy(img_dim=img_size,\n",
    "                                     lambda_illu=wavelength/1000,\n",
    "                                     ps=ps,\n",
    "                                     NA_obj=NA_obj,\n",
    "                                     NA_illu=NA_illu,\n",
    "                                     z_defocus=z_defocus,\n",
    "                                     chi=swing,\n",
    "                                     n_media=n_objective_media,\n",
    "                                     cali=True,\n",
    "                                     bg_option=bg_correction,\n",
    "                                     A_matrix=A,\n",
    "                                     QLIPP_birefringence_only=False,\n",
    "                                     pad_z=pad_z,\n",
    "                                     phase_deconv=mode,\n",
    "                                     illu_mode='BF',\n",
    "                                     use_gpu=use_gpu,\n",
    "                                     gpu_id=gpu_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e2ced99",
   "metadata": {},
   "source": [
    "### Reconstruct Raw Stokes images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "7f26bc4f-03d2-4b71-9553-f0b82d0176ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "S_raw = np.zeros((n_timepoints, 3, n_slices, *img_size), dtype='float')\n",
    "\n",
    "for t in range(n_timepoints):\n",
    "    S_raw_ = reconstructor.Stokes_recon(np.moveaxis(I_registered[t], 1, -1) - black_level)\n",
    "    S_raw[t] = np.moveaxis(S_raw_, -1, 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "956c86d8",
   "metadata": {},
   "source": [
    "#### View raw Stokes images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "02b29cc4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<Image layer 'S2_raw' at 0x7fae817d00d0>"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "viewer = napari.Viewer()\n",
    "viewer.add_image(S_raw[:,0], name='S0_raw', colormap='gray')\n",
    "viewer.add_image(S_raw[:,1], name='S1_raw', colormap='RdBu', visible=False)\n",
    "viewer.add_image(S_raw[:,2], name='S2_raw', colormap='RdBu', visible=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0be59db5",
   "metadata": {},
   "source": [
    "### Denoise S0 images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "1c1abe2c",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/ivan.ivanov/anaconda3/envs/mipolscope/lib/python3.9/site-packages/pywt/_multilevel.py:43: UserWarning: Level value of 2 is too high: all coefficients will experience boundary effects.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "S0_raw_denoised = wavelet_softThreshold(S_raw[:,0], 'db8', 200, level=2, axes=(1,2,3))\n",
    "S0_raw_denoised = np.expand_dims(S0_raw_denoised, axis=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de7c3fbb",
   "metadata": {},
   "source": [
    "#### Compare raw and denoised images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "64ba057f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<Image layer 'S0_raw_denoised' at 0x7fb74c5c2a30>"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "viewer = napari.Viewer()\n",
    "viewer.add_image(S_raw[:,0], name='S0_raw', colormap='gray')\n",
    "viewer.add_image(S0_raw_denoised[:,0], name='S0_raw_denoised', colormap='gray')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "91a4a767",
   "metadata": {},
   "source": [
    "### Normalize Stokes images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "25e70393",
   "metadata": {},
   "outputs": [],
   "source": [
    "S_norm = np.zeros_like(S_raw)\n",
    "\n",
    "for t in range(n_timepoints):\n",
    "    for z in range(n_slices):\n",
    "        S_norm[t,:,z] = reconstructor.Stokes_transform(np.concatenate([S0_raw_denoised[t,:,z], \n",
    "                                                                       S_raw[t,1:,z]]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88a3038e",
   "metadata": {},
   "source": [
    "#### View normalized Stokes images and background images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "507c0c6e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<Image layer 'S2_bg' at 0x7fb739f8a370>"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "viewer = napari.Viewer()\n",
    "viewer.add_image(S_norm[:,0], name='S0_norm', colormap='gray')\n",
    "viewer.add_image(S_norm[:,1], name='S1_norm', colormap='RdBu', visible=False, contrast_limits=(-0.1, 0.1))\n",
    "viewer.add_image(S_norm[:,2], name='S2_norm', colormap='RdBu', visible=False, contrast_limits=(-0.1, 0.1))\n",
    "\n",
    "viewer.add_image(S_bg[0], name='S0_bg', colormap='gray', visible=False, contrast_limits=viewer.layers[0].contrast_limits)\n",
    "viewer.add_image(S_bg[1], name='S1_bg', colormap='RdBu', visible=False, contrast_limits=(-0.1, 0.1))\n",
    "viewer.add_image(S_bg[2], name='S2_bg', colormap='RdBu', visible=False, contrast_limits=(-0.1, 0.1))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d4d6dc7",
   "metadata": {},
   "source": [
    "### Correct background"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "8a85dfe9",
   "metadata": {},
   "outputs": [],
   "source": [
    "S_corr = np.zeros_like(S_raw)\n",
    "\n",
    "for t in range(n_timepoints):\n",
    "    S_corr_ = reconstructor.Polscope_bg_correction(np.moveaxis(S_norm[t], 1, -1), S_bg)\n",
    "    S_corr[t] = np.moveaxis(S_corr_, -1, 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41806cd8",
   "metadata": {},
   "source": [
    "#### View background corrected Stokes images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "e93ef469",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<Image layer 'S2_corr' at 0x7fb74dd1b430>"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "viewer = napari.Viewer()\n",
    "viewer.add_image(S_corr[:,0], name='S0_corr', colormap='gray')\n",
    "viewer.add_image(S_corr[:,1], name='S1_corr', colormap='RdBu', visible=False, contrast_limits=(-0.02, 0.02))\n",
    "viewer.add_image(S_corr[:,2], name='S2_corr', colormap='RdBu', visible=False, contrast_limits=(-0.02, 0.02))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19578b98",
   "metadata": {},
   "source": [
    "### Denoise S1 and S2 images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "af1df934",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/ivan.ivanov/anaconda3/envs/mipolscope/lib/python3.9/site-packages/pywt/_multilevel.py:43: UserWarning: Level value of 2 is too high: all coefficients will experience boundary effects.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "# z projection\n",
    "z_chunk_size = 4\n",
    "n_z_chunks = n_slices//z_chunk_size\n",
    "S0_proj = np.zeros((n_timepoints, n_z_chunks, *img_size))\n",
    "S1_proj = np.zeros_like(S0_proj)\n",
    "S2_proj = np.zeros_like(S0_proj)\n",
    "\n",
    "for i in range(n_z_chunks):\n",
    "    S0_proj[:,i] = np.mean(S_corr[:, 0, i*z_chunk_size:(i+1)*z_chunk_size], axis=1)\n",
    "    S1_proj[:,i] = np.mean(S_corr[:, 1, i*z_chunk_size:(i+1)*z_chunk_size], axis=1)\n",
    "    S2_proj[:,i] = np.mean(S_corr[:, 2, i*z_chunk_size:(i+1)*z_chunk_size], axis=1)\n",
    "\n",
    "# wavelet denoising\n",
    "S1_denoised = wavelet_softThreshold(S1_proj, 'db8', 6e-3, level=2, axes=(1,2,3))\n",
    "S2_denoised = wavelet_softThreshold(S2_proj, 'db8', 6e-3, level=2, axes=(1,2,3))\n",
    "S_denoised = np.stack((S0_proj, S1_denoised, S2_denoised), axis=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "231a9ff5",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### View denoised Stokes images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "9590599c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<Image layer 'S2' at 0x7fb78ab103a0>"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "viewer = napari.Viewer()\n",
    "viewer.add_image(S_denoised[:,0], name='S0', colormap='gray')\n",
    "viewer.add_image(S_denoised[:,1], name='S1', colormap='RdBu', visible=False, contrast_limits=(-0.02, 0.02))\n",
    "viewer.add_image(S_denoised[:,2], name='S2', colormap='RdBu', visible=False, contrast_limits=(-0.02, 0.02))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94f3e706",
   "metadata": {},
   "source": [
    "## Compute transmission, phase, retardance, and orientation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "f7b32b1a",
   "metadata": {},
   "outputs": [],
   "source": [
    "retardance = np.zeros((n_timepoints, 1, n_z_chunks, *img_size))\n",
    "orientation = np.zeros_like(retardance)\n",
    "transmission = np.zeros_like(retardance)\n",
    "\n",
    "for t in range(n_timepoints):\n",
    "    phys_props_ = reconstructor.Polarization_recon(np.moveaxis(S_denoised[t], 1, -1))\n",
    "    retardance[t,0], orientation[t,0], transmission[t,0] = np.moveaxis(phys_props_, -1, 1)\n",
    "    \n",
    "transmission /= transmission[0].mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "c9fa9195",
   "metadata": {},
   "outputs": [],
   "source": [
    "phase = np.zeros((n_timepoints, 1, n_slices, *img_size))\n",
    "\n",
    "for t in range(n_timepoints):\n",
    "    phase_ = reconstructor.Phase_recon_3D(np.moveaxis(S_corr[t, 0], 0, -1), method='Tikhonov')\n",
    "    phase[t,0] = np.moveaxis(phase_,-1,0)\n",
    "    \n",
    "# project phase\n",
    "phase = np.mean(np.reshape(phase, (-1, 1, n_z_chunks, z_chunk_size, *img_size)), axis=-3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "becf20e4",
   "metadata": {},
   "source": [
    "# View reconstruction results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "01bd16e0",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<Image layer 'phase' at 0x7fb770731100>"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "viewer = napari.Viewer()\n",
    "viewer.add_image(retardance, name='retardance', colormap='gray', contrast_limits=(0, 0.02))\n",
    "viewer.add_image(orientation, name='orientation', colormap='gray', visible=False, contrast_limits=(0, np.pi))\n",
    "viewer.add_image(transmission, name='transmission', colormap='gray', visible=False, contrast_limits=(0.8, 1.2))\n",
    "viewer.add_image(phase, name='phase', colormap='gray', visible=False, contrast_limits=(-0.02, 0.02))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f339898c",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}