Add notebook demo for local pipeline

examples/local_pipeline_tutorial.ipynb

@@ -0,0 +1,367 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "2be64610-5160-415b-a59d-92966bc7f8b6",
+   "metadata": {},
+   "source": [
+    "Run inference and segmentation tasks locally \n",
+    "==================================\n",
+    "This tutorial goes through steps to run inference and segmentation tasks with SEURON deployed locally by docker compose. **You should only use this notebook in JupyterLab hosted by a local SEURON deployment**. In order to run all the examples, you need a recent Linux system with **compatible GPU toolkit**. Make sure you follow the steps in the README file to create the local deployment.\n",
+    "\n",
+    "The jupyterlab directory is mounted from a docker volume so the notebooks stored here are preserved until you delete the docker volume itself. A read-only copy of this tutorial is stored in the examples directory, in case you want to reverse changes made."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "38d06c0b-ef96-4ec0-a6f0-509467c5480b",
+   "metadata": {},
+   "source": [
+    "Load seuronbot extension\n",
+    "========================\n",
+    "In order to communicate with SEURON in a notebook, we need to load the seuronbot extension. **The seuronbot extension can only be loaded by one notebook, trying to load another instance of the extension will result in an error. If you want to use the extension in a different notebook, make sure you reset the kernel of the current notebook to unload the extension**"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "928c2fc4-f5a2-41a0-967f-2e5545039edc",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "%load_ext seuronbot_ext"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "975ca35d-06e5-4423-a348-80ee70f59edd",
+   "metadata": {},
+   "source": [
+    "Print a list of commands seuronbot supports\n",
+    "===========================================\n",
+    "In this tutorial we use jupyter frontend to demonstrate seuronbot commands, which are always initiated by the line/cell magic `seuronbot`. Depending on whether the command requires additional inputs, either the **line magic** `%seuronbot` or **cell magic** `%%seuronbot` should be used.\n",
+    "\n",
+    "When using the **slack frontend** instead, one should replace `%seuronbot` and `%%seuronbot` by the bot's name. The name of the bot has the form of `seuron-worker-{DEPLOYMENT}`. For Google cloud deployment, `DEPLOYMENT` is the name you choose for your deployment, for local deployment, a randomly generated `DEPLOYMENT` is used and write to `.env.local` file. The bot will send a hello message in the notification channel you specified when setting up the deployment."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a7f479fa-155c-4fb0-bb8c-fc36d5ad46cc",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "%seuronbot help"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "5d6ef176-445a-4cb0-b973-1533a5b7451d",
+   "metadata": {
+    "tags": []
+   },
+   "source": [
+    "Playground: a small cutout from FlyWire\n",
+    "=====================================\n",
+    "In order to demonstrate the pipeline, we use a small cutout from FAFB v14, which is the image stack FlyWire is based on. The bounding box of the cutout is [38000, 11500, 3530, 38512, 12012, 3730] at 16nm x 16nm x 40nm voxel size. In the cell below, The cutout are visualized using neuroglancer."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f983bcf8-f07a-44cb-a622-a15d11497558",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "from IPython.display import IFrame, HTML\n",
+    "\n",
+    "ng_link = \"https://cj-find-path-test-dot-neuroglancer-dot-seung-lab.appspot.com/#!%7B%22dimensions%22:%7B%22x%22:%5B1.6e-8%2C%22m%22%5D%2C%22y%22:%5B1.6e-8%2C%22m%22%5D%2C%22z%22:%5B4e-8%2C%22m%22%5D%7D%2C%22position%22:%5B38266.5%2C11729.5%2C3616.5%5D%2C%22crossSectionScale%22:0.9999999999999981%2C%22projectionOrientation%22:%5B-0.24463512003421783%2C0.09365575760602951%2C-0.11994557827711105%2C-0.9575987458229065%5D%2C%22projectionScale%22:2157.378810418306%2C%22layers%22:%5B%7B%22type%22:%22image%22%2C%22source%22:%22precomputed://https://bossdb-open-data.s3.amazonaws.com/flywire/fafbv14%22%2C%22tab%22:%22source%22%2C%22name%22:%22fafbv14%22%7D%2C%7B%22type%22:%22annotation%22%2C%22source%22:%7B%22url%22:%22local://annotations%22%2C%22transform%22:%7B%22outputDimensions%22:%7B%22x%22:%5B1.6e-8%2C%22m%22%5D%2C%22y%22:%5B1.6e-8%2C%22m%22%5D%2C%22z%22:%5B4e-8%2C%22m%22%5D%7D%7D%7D%2C%22tool%22:%22annotateBoundingBox%22%2C%22tab%22:%22annotations%22%2C%22annotations%22:%5B%7B%22pointA%22:%5B38000%2C11500%2C3530%5D%2C%22pointB%22:%5B38512%2C12012%2C3730%5D%2C%22type%22:%22axis_aligned_bounding_box%22%2C%22id%22:%2215fc1f9ae16810f0dd8aa4db2ac5eceed4ff7f56%22%7D%5D%2C%22name%22:%22annotation%22%7D%5D%2C%22selectedLayer%22:%7B%22layer%22:%22annotation%22%7D%2C%22layout%22:%22xy-3d%22%7D\"\n",
+    "display(IFrame(src=ng_link, width=1200, height=600))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "0016bcbd-9217-4bfd-ab63-35547847bf03",
+   "metadata": {},
+   "source": [
+    "Run inference to create affinity map\n",
+    "====================================\n",
+    "To segment the cutout we picked, we first need to generate affinity map for it. The procedure is divided into two steps: First one submits the parameters describing the tasks, using the `update parameters` command. The bot performs sanity checks to prevent common mistakes. After the sanity check succeeds. One can launch the pipeline with `run pipeline` commands.\n",
+    "\n",
+    "Submit inference parameters\n",
+    "--------------------------------------------------------\n",
+    "The command will send the python code in the current cell to SEURON. The cell content must be **self-contained**, no reference to variables or modules from other cells. The cell have to define a `submit_parameter` function requiring no arguments. The function should return a dictionary containing the inference parameters or a list of dictionaries for several tasks. The bot will execute this function in a separate environment and use the returned parameters for sanity checks.\n",
+    "\n",
+    "SEURON uses [the zettaai fork of chunkflow](https://github.com/ZettaAI/chunkflow) for inference tasks, `ranlu/chunkflow:zettaai` pointing to the latest build from the zettaiai. The parameters for each run can be divided into roughly two types: parameters for IO and parameters for the inference model. The IO parameters contain information of the input images: the location, resolution and bounding box, the output will be stored in `OUTPUT_PATH`, or `{OUTPUT_PREFIX}{NAME}` when `OUTPUT_PATH` is not defined. \n",
+    "The parameters for the inference model must be adjusted for each model, the parameters here are suitable for [the FlyWire model](https://github.com/seung-lab/DeepEM/releases/tag/0.0.6).\n",
+    "\n",
+    "For local processes, we store the final and intermediate results to `/tmp`, which is automatically mounted from the host for all containers managed by SEURON. If you want to use some other volume, make sure you add them to the deployment file and have proper permission."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2e4d8fe4-be0f-40bd-8785-9fb0657e1d7f",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "%%seuronbot update parameters\n",
+    "def submit_inference_parameters():\n",
+    "    bbox = [38000, 11500, 3530, 38512, 12012, 3730]\n",
+    "    io_param = {\n",
+    "        \"IMAGE_PATH\": \"https://bossdb-open-data.s3.amazonaws.com/flywire/fafbv14\",\n",
+    "        \"IMAGE_FILL_MISSING\": True,\n",
+    "        \"IMAGE_RESOLUTION\": [16, 16, 40],\n",
+    "        \"BBOX\": bbox,\n",
+    "        \"OUTPUT_PREFIX\": \"file:///tmp/scratch/ng/test_aff/\",\n",
+    "    }\n",
+    "    aff_param = {\n",
+    "        \"INPUT_PATCH_SIZE\": [128,128,20],\n",
+    "        \"OUTPUT_PATCH_SIZE\": [96,96,16],\n",
+    "        \"OUTPUT_CHANNELS\": 3,\n",
+    "        \"INPUT_PATCH_OVERLAP_RATIO\": 0.5,\n",
+    "        \"CHUNKFLOW_IMAGE\": \"ranlu/chunkflow:zettaai\",\n",
+    "        \"ONNX_MODEL_PATH\": \"https://github.com/seung-lab/DeepEM/releases/download/0.0.6/flywire_v0.2.onnx\",\n",
+    "    }\n",
+    "    return [\n",
+    "                {\n",
+    "                    \"NAME\": \"aff_local_test\",\n",
+    "                    **io_param,\n",
+    "                    **aff_param,\n",
+    "                },\n",
+    "           ]\n",
+    "\n",
+    "def submit_parameters():\n",
+    "    inf_param = submit_inference_parameters()\n",
+    "    return inf_param"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "0914280f-0e58-4095-9ef8-6495deb38455",
+   "metadata": {},
+   "source": [
+    "Launch inference run\n",
+    "--------------------\n",
+    "Once the sanity check succeeded, we can start the inference using `%seuronbot run pipeline` command. The first message after the run is triggered is the cancel token to use if you want to cancel the run."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "304cd6f4-18d0-49d5-84ea-cee6193102f7",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "%seuronbot run pipeline"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "80952c5a-d904-4e06-a211-9957904cf740",
+   "metadata": {},
+   "source": [
+    "Segment the affinity map\n",
+    "========================\n",
+    "After the affinity map is generated, we use it to create segmentation of the Flywire cutout. The procedure is similar to the inference tasks, first use `update parameter` command to submit the parameters describing the tasks, then use `run pipeline` to trigger the run.\n",
+    "\n",
+    "Submit segmentation parameters\n",
+    "----------------------------------\n",
+    "For segmentation tasks SEURON uses [ABISS](https://github.com/seung-lab/abiss), `ranlu/abiss:main` is built from the main branch. The command creates both flat segmentations and inputs for [PyChunkedGraph](https://github.com/seung-lab/PyChunkedGraph). The segmentation requires affinity map as an input, make sure you update the path to the affinity map if you modified it in the previous step."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2a2b7b05-465c-4344-bc6b-4ae952ff8e71",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "%%seuronbot update parameters\n",
+    "from datetime import datetime\n",
+    "\n",
+    "def submit_segmentation_parameters(aff_path):\n",
+    "    bbox = [38000, 11500, 3530, 38512, 12012, 3730]\n",
+    "    now = datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
+    "    io_param = {\n",
+    "        \"SCRATCH_PREFIX\": \"file:///tmp/scratch/\",\n",
+    "        \"IMAGE_PATH\": \"https://bossdb-open-data.s3.amazonaws.com/flywire/fafbv14\",\n",
+    "        \"AFF_PATH\": aff_path,\n",
+    "        \"BBOX\": bbox,\n",
+    "        \"NG_PREFIX\": \"file:///tmp/scratch/ng/\",\n",
+    "    }\n",
+    "    seg_param = {\n",
+    "        \"WS_HIGH_THRESHOLD\": \"0.99\",\n",
+    "        \"WS_LOW_THRESHOLD\": \"0.01\",\n",
+    "        \"WS_SIZE_THRESHOLD\": \"200\",\n",
+    "        \"AGG_THRESHOLD\": \"0.2\",\n",
+    "        \"WORKER_IMAGE\": \"ranlu/abiss:main\",\n",
+    "        \"SKIP_SKELETON\": True,\n",
+    "    }\n",
+    "    return [\n",
+    "                {\n",
+    "                    \"NAME\": f\"seg_fp32_{now}\",\n",
+    "                    **io_param,\n",
+    "                    **seg_param,\n",
+    "                },\n",
+    "           ]\n",
+    "\n",
+    "def submit_parameters():\n",
+    "    aff_path = \"file:///tmp/scratch/ng/test_aff/aff_local_test\"\n",
+    "    run_param = submit_segmentation_parameters(aff_path)\n",
+    "    return run_param"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "86742989-1771-40d7-baef-0b5642b1ef9a",
+   "metadata": {},
+   "source": [
+    "Launch segmentation run\n",
+    "-----------------------\n",
+    "Same as inference, we can start the segmentation using `%seuronbot run pipeline` command. And you will receive the cancel token at the starting of the run"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e2f26ea5-a287-439b-b7bb-e7c98edb41ac",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "%seuronbot run pipeline"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "31d224e0-e468-42d5-bcd9-be0db166d707",
+   "metadata": {},
+   "source": [
+    "Combine runs\n",
+    "============\n",
+    "Instead of running single inference run or segmentation run, one can send parameters of multiple runs with a single `update parameters` command or even combine inference and segmentation. For example we can combine the previous two runs together"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e1f86a04-8449-4a4b-9a0f-6cc3487e3c2f",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "%%seuronbot update parameters\n",
+    "from datetime import datetime\n",
+    "\n",
+    "bbox = [38000, 11500, 3530, 38512, 12012, 3730]\n",
+    "def submit_segmentation_parameters(aff_path):\n",
+    "    now = datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
+    "    io_param = {\n",
+    "        \"SCRATCH_PREFIX\": \"file:///tmp/scratch/\",\n",
+    "        \"IMAGE_PATH\": \"https://storage.googleapis.com/neuroglancer-fafb-data/fafb_v14/fafb_v14_clahe\",\n",
+    "        \"AFF_PATH\": aff_path,\n",
+    "        \"BBOX\": bbox,\n",
+    "        \"NG_PREFIX\": \"file:///tmp/scratch/ng/\",\n",
+    "    }\n",
+    "    seg_param = {\n",
+    "        \"WS_HIGH_THRESHOLD\": \"0.99\",\n",
+    "        \"WS_LOW_THRESHOLD\": \"0.01\",\n",
+    "        \"WS_SIZE_THRESHOLD\": \"200\",\n",
+    "        \"AGG_THRESHOLD\": \"0.2\",\n",
+    "        \"WORKER_IMAGE\": \"ranlu/abiss:main\",\n",
+    "        \"SKIP_SKELETON\": True,\n",
+    "    }\n",
+    "    return [\n",
+    "                {\n",
+    "                    \"NAME\": f\"seg_fp32_{now}\",\n",
+    "                    **io_param,\n",
+    "                    **seg_param,\n",
+    "                },\n",
+    "           ]\n",
+    "\n",
+    "\n",
+    "def submit_inference_parameters():\n",
+    "    now = datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
+    "    io_param = {\n",
+    "        \"IMAGE_PATH\": \"https://storage.googleapis.com/neuroglancer-fafb-data/fafb_v14/fafb_v14_clahe\",\n",
+    "        \"IMAGE_FILL_MISSING\": True,\n",
+    "        \"IMAGE_RESOLUTION\": [16, 16, 40],\n",
+    "        \"BBOX\": bbox,\n",
+    "        \"OUTPUT_PREFIX\": \"file:///tmp/scratch/ng/test_aff/\",\n",
+    "    }\n",
+    "    aff_param = {\n",
+    "        \"INPUT_PATCH_SIZE\": [128,128,20],\n",
+    "        \"OUTPUT_PATCH_SIZE\": [96,96,16],\n",
+    "        \"OUTPUT_CHANNELS\": 3,\n",
+    "        \"INPUT_PATCH_OVERLAP_RATIO\": 0.5,\n",
+    "        \"CHUNKFLOW_IMAGE\": \"ranlu/chunkflow:ort\",\n",
+    "        \"ONNX_MODEL_PATH\": \"https://github.com/seung-lab/DeepEM/releases/download/0.0.6/flywire_v0.2.onnx\",\n",
+    "    }\n",
+    "    return [\n",
+    "                {\n",
+    "                    \"NAME\": f\"aff_{now}\",\n",
+    "                    **io_param,\n",
+    "                    **aff_param,\n",
+    "                },\n",
+    "           ]\n",
+    "\n",
+    "def submit_parameters():\n",
+    "    inf_param = submit_inference_parameters()\n",
+    "    aff_path = inf_param[0]['OUTPUT_PREFIX']+inf_param[0][\"NAME\"]\n",
+    "    seg_param = submit_segmentation_parameters(aff_path)\n",
+    "    return inf_param+seg_param"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d118d34e-e029-4dab-b230-f86684476d6a",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "%seuronbot run pipeline"
+   ]
+  }
+ ],
+ "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.10.12"
+  },
+  "toc-autonumbering": false,
+  "toc-showcode": true,
+  "toc-showmarkdowntxt": false
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}