Applying Richardson-Lucy Deconvolution from ITK on AOLLSM using Dask

AOLLSM_ITK.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Import everything"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib notebook"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import importlib\n",
    "import os\n",
    "\n",
    "import numpy\n",
    "import numpy as np\n",
    "\n",
    "import dask\n",
    "import dask.array\n",
    "import dask.array as da\n",
    "import dask.array.image\n",
    "\n",
    "import zarr"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setup Dask Cluster"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import dask\n",
    "import dask.config\n",
    "\n",
    "from dask.distributed import Client, LocalCluster\n",
    "\n",
    "# Use nbserverproxy (now jupyter-server-proxy to proxy status page)\n",
    "dask.config.set({\"distributed.dashboard.link\": \"/proxy/{port}/status\"})\n",
    "\n",
    "cluster = LocalCluster()\n",
    "client = Client(cluster)\n",
    "client"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load our data from last time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "z = zarr.open(\"/public/AOLLSMData_m4_raw.zarr/\")\n",
    "z.tree()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "imgs = da.from_zarr(\"/public/AOLLSMData_m4_raw.zarr/\", \"data\")\n",
    "imgs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "imgs = da.from_zarr(\"/public/AOLLSMData_m4_raw.zarr/\", \"data\")\n",
    "imgs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load our Point Spread Function (PSF)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls /public/AOLLSMData/m4/psfs_z0p1/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dirpath = \"/public/AOLLSMData/m4/psfs_z0p1/\"\n",
    "\n",
    "psf = []\n",
    "for fn in sorted(os.listdir(dirpath)):\n",
    "    fn = os.path.join(dirpath, fn)\n",
    "    psf.append(dask.array.image.imread(fn))\n",
    "psf = da.stack(psf)\n",
    "\n",
    "psf"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Convert data to float32 for computation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "imgs = imgs.astype(np.float32)\n",
    "psf = psf.astype(np.float32)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load ITK on workers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_itk():\n",
    "    mod = importlib.import_module(\"itk\")\n",
    "    getattr(mod, \"PyBuffer\")\n",
    "    getattr(mod, \"Image\")\n",
    "    getattr(mod, \"RichardsonLucyDeconvolutionImageFilter\")\n",
    "\n",
    "client.run(load_itk);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Apply Richardson-Lucy Deconvolution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def itk_RichardsonLucyDeconvolutionImageFilter(img, psf):\n",
    "    # The module does not pickle\n",
    "    import itk\n",
    "    \n",
    "    # Drop first two singleton axes (need 3-D array at most)\n",
    "    img = img[0, 0]\n",
    "    psf = psf[0, 0]\n",
    "\n",
    "    # Get ITK Image buffers from NumPy arrays (assumes 3-D float32 arrays)\n",
    "    img = itk.PyBuffer[itk.Image.F3].GetImageFromArray(img)\n",
    "    psf = itk.PyBuffer[itk.Image.F3].GetImageFromArray(psf)\n",
    "    \n",
    "    # Configure filter and get result reference\n",
    "    op = itk.RichardsonLucyDeconvolutionImageFilter[itk.Image.F3, itk.Image.F3].New()\n",
    "    op.SetInput(img)\n",
    "    op.SetKernelImage(psf)\n",
    "    r = op.GetOutput()\n",
    "\n",
    "    # Compute result and view as NumPy array\n",
    "    r.Update()\n",
    "    r = itk.PyBuffer[itk.Image.F3].GetArrayViewFromImage(r)\n",
    "    \n",
    "    # Readd singleton axes\n",
    "    r = r[None, None]\n",
    "\n",
    "    return r"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "imgs_deconv = da.map_blocks(itk_RichardsonLucyDeconvolutionImageFilter, imgs, psf)\n",
    "imgs_deconv"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Compute and write to disk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "imgs_deconv.to_zarr(\"/public/AOLLSMData_m4_raw.zarr/\", \"deconvolved\", overwrite=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}