{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load Zarr Image from a public S3 store and analyze it in parallel\n",
    "\n",
    "The images are taken from  the paper \"In Toto Imaging and Reconstruction of Post-Implantation Mouse Development at the Single-Cell Level\" published October 2018 in Cell: https://doi.org/10.1016/j.cell.2018.09.031. The images can be viewed online in the [Image Data Resource](http://idr.openmicroscopy.org/webclient/?show=project-502).\n",
    "\n",
    "Some images from the Image Data Resource (IDR) have been converted into Zarr using [bioformats2raw](https://github.com/glencoesoftware/bioformats2raw) and stored on a public S3 store.\n",
    "\n",
    "What is **Zarr**? [Zarr](https://zarr.readthedocs.io/en/stable/) is a young format for the storage of chunked, compressed, N-dimensional arrays.\n",
    "\n",
    "We will use [Dask](https://docs.dask.org/en/latest/), a flexible library for parallel computing in Python. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Install dependencies if required\n",
    "The cell below will install dependencies if you choose to run the notebook in [Google Colab](https://colab.research.google.com/notebooks/intro.ipynb#recent=true). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install omero-py dask_image zarr fsspec>=0.3.3 aiohttp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Import "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import dask\n",
    "import dask.array as da\n",
    "import dask_image.ndfilters\n",
    "import dask_image.ndmeasure\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ID of the IDR image to analyze\n",
    "The original image ID is used to identify the file on the S3 store."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "image_id = 4007801"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load Zarr Image from S3\n",
    "\n",
    "The method will return a dask array without any binary data. The dimension order is ``(TCZYX)``. Data will be loaded when requested during the analysis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_binary_from_s3(id, resolution='4'):\n",
    "    endpoint_url = 'https://uk1s3.embassy.ebi.ac.uk/'\n",
    "    root = 'idr/zarr/v0.1/%s.zarr/%s/' % (id, resolution)\n",
    "    return da.from_zarr(endpoint_url + root)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the dask array\n",
    "data = load_binary_from_s3(image_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(532, 2, 988, 128, 135)\n"
     ]
    }
   ],
   "source": [
    "# Check the dimension of the array\n",
    "print(data.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Segment the image\n",
    "This could be replaced by the analysis you wish to perform on a plane."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def analyze(t, c, z):\n",
    "    plane = data[t, c, z, :, :] \n",
    "    smoothed_image = dask_image.ndfilters.gaussian_filter(plane, sigma=[1, 1])\n",
    "    threshold_value = 0.33 * da.max(smoothed_image).compute()\n",
    "    threshold_image = smoothed_image > threshold_value\n",
    "    label_image, num_labels = dask_image.ndmeasure.label(threshold_image)\n",
    "    name = \"t:%s, c: %s, z:%s\" % (t, c, z)\n",
    "    return label_image, name"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Use Dask Delayed "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We use ``dask.delayed`` on our function so it records what we want to compute as a task into a graph that will run later on parallel hardware.\n",
    "\n",
    "Due to the size of the image, we only analyze in the context of this notebook a few planes around the middle z-section and middle timepoint for the first channel. In total, we are analyzing 100 planes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Delayed('analyze-3a88bfbf-a462-4ece-8864-fe08e218a769'), Delayed('analyze-fdfb0923-c9b2-46bb-a047-de00766157c1'), Delayed('analyze-fad31d2e-e4d5-42f7-a296-8b33275b0eac'), Delayed('analyze-1150e10e-efea-432b-8737-a0b3f59175f3'), Delayed('analyze-94d6c7ec-ac02-44ea-b8df-dd3632f2efcf'), Delayed('analyze-b8cd0a5b-b194-4b29-9a6d-55f57f10d2d7'), Delayed('analyze-600c6ff3-4357-49f7-a5f8-7847f60c0cc9'), Delayed('analyze-e46abfd1-4a7b-45e8-9234-63170e4f3c87'), Delayed('analyze-25f4e1ff-56f8-4d60-a447-2b38ea48cbd4'), Delayed('analyze-78b9c9a7-90f0-421d-ba12-b4ed6323658b'), Delayed('analyze-49e09a50-e214-416b-8d80-787ec41cb8ee'), Delayed('analyze-908c97da-9529-4f96-9274-09dd345164a8'), Delayed('analyze-6dcdbde7-f7e2-41dd-b086-b2fdd530a188'), Delayed('analyze-8c4c4fba-1dec-4331-8797-a6c979243b8f'), Delayed('analyze-0b7d2c19-e855-4cea-acd0-05322463f6cf'), Delayed('analyze-dd7e3627-7187-4561-b291-e3783400f147'), Delayed('analyze-a0a72496-d606-481f-921d-ff6fbf089402'), Delayed('analyze-50b0892b-b308-4469-9189-a4366f1ff59b'), Delayed('analyze-06603f0d-a8f0-48e5-a8bb-d90a048ec4e4'), Delayed('analyze-aca3a562-0874-47ed-85be-fc8fe9cf6e45'), Delayed('analyze-d5bb9fce-e84d-455d-b151-4786dbc155b3'), Delayed('analyze-a1312d58-5f1d-4407-a589-be0154dde55c'), Delayed('analyze-098f9bc5-f035-4d59-9fa9-4b2e4ec57009'), Delayed('analyze-dd5ea0f2-a9fc-4690-ad7d-a4ab9e678f37'), Delayed('analyze-425cf197-ba26-4226-8cd5-cd14e8d46e83'), Delayed('analyze-75485845-5c46-4fce-9e1c-27fa43a9c9df'), Delayed('analyze-cad1dc73-9c16-4090-af76-a9bb681c33e1'), Delayed('analyze-b88d95e0-d588-4947-98c0-9d322ddebc4c'), Delayed('analyze-7c63856d-b3dc-40cd-9d36-7bd7b2eed23e'), Delayed('analyze-6c36041b-6abc-4c38-8611-f0e9b93e8c38'), Delayed('analyze-c0738f22-fb9d-4043-bad3-0b11bfe1524c'), Delayed('analyze-abd86334-007a-4878-80d2-f788f9dfab1c'), Delayed('analyze-94e66bd4-03e0-4e4a-9d60-913eb21cbd22'), Delayed('analyze-37f07d5c-5078-427f-ab01-3ef0a669d7d8'), Delayed('analyze-5fbc1103-46bd-4f1c-a16d-b18889296583'), Delayed('analyze-8140c327-0d3f-4e81-b66a-93b8615c6240'), Delayed('analyze-bf39c50e-a068-4c64-94fe-f3d2d38a5807'), Delayed('analyze-e4255387-9d1e-4a20-8099-9d6184e594a2'), Delayed('analyze-f6ffdaf6-4e29-4d4d-a471-e81f8c0899bb'), Delayed('analyze-20cfcd6d-c0de-4545-99bd-62ad9cf6af6e'), Delayed('analyze-4c338a13-a710-4640-85c9-fcebf5abee86'), Delayed('analyze-4bb403be-bc90-4bac-b825-d0ccf4755bc7'), Delayed('analyze-65d6a51d-63de-476f-972c-6a83dcb5716a'), Delayed('analyze-4821f972-ffed-4efa-848b-5c7bab72ffe6'), Delayed('analyze-145c8135-7644-4277-9f19-a5059b98ca89'), Delayed('analyze-99408488-9627-48d5-869b-d03b41bb10c2'), Delayed('analyze-c5c5fa7d-0139-4efb-8e14-bcdf817a3c63'), Delayed('analyze-0d9b7ff5-3306-4a99-9173-af7e104a1c85'), Delayed('analyze-14bcf3c1-7b00-40b3-8fa8-4f2ff5e61f01'), Delayed('analyze-d1cbcfea-0051-43a8-aba5-ad8616c1ca0e'), Delayed('analyze-0de60a29-8c38-4381-a47a-5ca37aaded1e'), Delayed('analyze-72841ccd-b743-4b04-adb3-1bfc810dfe23'), Delayed('analyze-57240eec-0736-4d0a-a9d2-1e52490e9034'), Delayed('analyze-3ea23a7e-9334-4490-a491-7ea12264e207'), Delayed('analyze-64f513b9-fc6f-422b-827e-6a523989a31e'), Delayed('analyze-5e8953fc-f451-4baf-bf53-1a86d7aaf7c8'), Delayed('analyze-5f8ece6a-f542-4c8d-8642-047b367a4c1e'), Delayed('analyze-ad9e9aa8-22d1-463e-929a-7084b2a3e8f2'), Delayed('analyze-b1964c8e-de90-4858-9542-1561576e3f00'), Delayed('analyze-3df0b505-8b7f-4fbe-aeb8-360df5e8f1bb'), Delayed('analyze-075ef5dd-7bab-4248-8f66-7042a29d7492'), Delayed('analyze-91d0ba91-a9c5-47da-afe8-ef7c02c4d7b2'), Delayed('analyze-98edfa2a-c86a-444a-8278-1f9bf6086edd'), Delayed('analyze-7124e400-fe16-4f4a-99ac-a7d8831bd5b0'), Delayed('analyze-457c2e2a-06be-4f59-a471-11500fe9a359'), Delayed('analyze-d14c25bd-7b39-449b-81c5-4e279f9e6658'), Delayed('analyze-fe33ed2d-4cf9-4fd5-afad-6fbe283f1348'), Delayed('analyze-e5e158fa-9ce6-4e63-a2e2-65e19088e4da'), Delayed('analyze-b70d3300-3f84-430a-be5c-a7e92b27a812'), Delayed('analyze-f090696a-aceb-47a5-a63f-fdb607577b90'), Delayed('analyze-0d8466fe-3fd9-43dc-bda0-e32774904cab'), Delayed('analyze-351bb272-b30c-4e71-a044-8ebe7198aac0'), Delayed('analyze-46a52fc4-5313-459a-9f56-c6f72a1ee3a2'), Delayed('analyze-08cd7344-1682-40d8-a2d1-f637d8fe4ba3'), Delayed('analyze-fd5ab082-0d5f-4fbf-a924-69ec98023912'), Delayed('analyze-4677672e-81a0-434e-9047-62ea12fa930e'), Delayed('analyze-a284e3d2-1c47-4a0d-bd38-e63ea9c06412'), Delayed('analyze-5d7a8afe-b04e-410d-8bb2-bc026ce1f915'), Delayed('analyze-c1b92548-1d76-4c57-b024-f6ac000e52ba'), Delayed('analyze-3bb5200e-206b-4588-8b17-f0fb8f2d6613'), Delayed('analyze-f4de4d6d-9f8b-45f4-8000-4b2b71c76230'), Delayed('analyze-84669fe1-d960-4ef1-9e7b-8a2c9fd501b9'), Delayed('analyze-10b6d81a-4a9e-404c-9d2c-f95d8d55d479'), Delayed('analyze-e5e31017-1a0d-4803-b469-839ed9745599'), Delayed('analyze-c87957c1-8e09-411f-9555-37d90cc849af'), Delayed('analyze-b3959e90-4ad2-40da-bbbc-91e5fc434408'), Delayed('analyze-0f74d458-174d-41c6-beb4-9407b9315f1b'), Delayed('analyze-4da4393c-b5ce-48d2-abce-a27f5d980dbe'), Delayed('analyze-d180bbce-6158-4549-bd9d-7a70ac57c8b6'), Delayed('analyze-2705e51b-302e-4256-9ae2-b553ae3daf3e'), Delayed('analyze-e007bf34-3fa7-4afc-88a0-4d90de06d097'), Delayed('analyze-94ba240f-234c-4226-af4d-359ba9996ae6'), Delayed('analyze-5a58da0b-0e7c-416e-b3fd-3c9fae1e0c51'), Delayed('analyze-94010098-f6d4-46a2-a40c-842df7ac3ecb'), Delayed('analyze-35e5dd15-8341-4547-9540-71c6f866755b'), Delayed('analyze-8473b151-37eb-4f04-a120-18ba8be04150'), Delayed('analyze-7283f242-b87e-4f04-81df-875f5fc200b2'), Delayed('analyze-3c99e03c-34c7-4f24-ac64-1e504f5afeec'), Delayed('analyze-4755fc24-5b8d-4b26-8ec4-aba590b717b1'), Delayed('analyze-33460b40-b8f4-415e-83b1-bbf013f0e6e7')]\n"
     ]
    }
   ],
   "source": [
    "lazy_results = []\n",
    "middle_z = data.shape[2] // 2\n",
    "middle_t = data.shape[0] // 2\n",
    "range_t = 5\n",
    "range_z = 5\n",
    "range_c = 1\n",
    "for t in range(middle_t - range_t, middle_t + range_t):\n",
    "    for z in range(middle_z - range_z, middle_z + range_z):\n",
    "        for c in range(range_c):\n",
    "            lazy_result = dask.delayed(analyze)(t, c, z)\n",
    "            lazy_results.append(lazy_result)\n",
    "print(lazy_results)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run in parallel"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The lazy_results list contains information about ``number_t*number_z*number_c`` calls to our function ``analyze`` that have not yet been run. We call ``dask.compute`` when we want the results.\n",
    "The binary data is loaded from the S3 store during the ``compute`` phase and the analysis is performed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1min 42s, sys: 4.93 s, total: 1min 47s\n",
      "Wall time: 1min 48s\n"
     ]
    }
   ],
   "source": [
    "%time results = dask.compute(*lazy_results)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Display the results\n",
    "\n",
    "Display the segmented planes. Use the slider to select the plane."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "458471cce9224dd0877828beaffaa2f4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(IntSlider(value=0, continuous_update=False, description='Select Plane', max=99), Output(…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<function __main__.display(i=0)>"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from ipywidgets import *\n",
    "\n",
    "def display(i=0):\n",
    "    r, name = results[i]\n",
    "    fig = plt.figure(figsize=(10, 10))\n",
    "    plt.subplot(121)\n",
    "    plt.imshow(r)\n",
    "    plt.title(name)\n",
    "    fig.canvas.flush_events()\n",
    "\n",
    "interact(display, i= widgets.IntSlider(value=0, min=0, max=len(results)-1, step=1, description=\"Select Plane\", continuous_update=False))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![Overview](./images/idr0044_segmented.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "### License (BSD 2-Clause)\n",
    "Copyright (C) 2020-2022 University of Dundee. All Rights Reserved.\n",
    "\n",
    "Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:\n",
    "\n",
    "Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.\n",
    "Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.\n",
    "THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}