{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Download benchmarking data from S3 with Neuroglancer\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### This notebook explains how to:\n",
    "(1) Read benchmarking data from S3 via Neuroglancer \n",
    "\n",
    "(2) download raw benchmarking data to your local computer \n",
    "\n",
    "#### Quick notes on the benchmarking data:\n",
    "\n",
    "In octree format, data is labled in folders, labeled test_1 through test_25 and validation_1 through validation_25. \n",
    "\n",
    "If when downloading, you get a reshape error, try first uploading segments and then re-uploading the volumes. \n",
    "\n",
    "Known issues with a few of the files: \n",
    "\n",
    "- test_9,test_10 - didnt seem to have good swc alignment\n",
    "\n",
    "- test_24 - issues with the image\n",
    "\n",
    "- validation_11 - seems to be a shift between swcs and the image\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import napari\n",
    "from napari.utils import nbscreenshot"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Define locations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "c:\\users\\shrey\\anaconda3\\envs\\ndd-windows\\lib\\site-packages\\python_jsonschema_objects\\__init__.py:53: UserWarning: Schema version http://json-schema.org/draft-04/schema not recognized. Some keywords and features may not be supported.\n",
      "  self.schema[\"$schema\"]\n"
     ]
    }
   ],
   "source": [
    "from brainlit.utils import session\n",
    "from brainlit.utils.Neuron_trace import NeuronTrace\n",
    "\n",
    "# #Can change to test_\"1-25\", validation_\"1-25\"\n",
    "dest = \"s3://open-neurodata/brainlit/benchmarking_data/validation_7\"\n",
    "dest_segments = \"s3://open-neurodata/brainlit/benchmarking_data/validation_7\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Create Neuroglancer session & download benchmarking volume"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "sess = session.NeuroglancerSession(\n",
    "    url=dest, url_segments=dest_segments, mip=0\n",
    ")  # create session object\n",
    "img, bounds, vertices = sess.pull_vertex_list(\n",
    "    1, [1], 0, expand=True\n",
    ")  # get full benchmarking image"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Download a specific .swc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Downloading: 100%|████████████████████████████████████| 1/1 [00:00<00:00,  6.65it/s]\n",
      "Downloading: 100%|████████████████████████████████████| 1/1 [00:00<00:00,  6.07it/s]\n"
     ]
    }
   ],
   "source": [
    "seg_id = 1  # Can change\n",
    "\n",
    "G_paths = sess.get_segments(seg_id, bounds, rounding=False)\n",
    "G = G_paths[0]\n",
    "paths = G_paths[1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Visualize with napari"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# viewer = napari.Viewer(ndisplay=3)\n",
    "# viewer.add_image(img)\n",
    "# viewer.add_shapes(data=paths, shape_type='path', edge_width=1.0, edge_color='blue', opacity=0.8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# nbscreenshot(viewer, canvas_only = True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Download raw benchmarking data\n",
    "\n",
    "###### This will download the benchmarking data in .tif and .swc format to a local destination"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "from botocore import UNSIGNED\n",
    "from botocore.client import Config\n",
    "import os\n",
    "from pathlib import Path\n",
    "import numpy as np\n",
    "from skimage import io\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create directories"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading segments to C:\\Users\\shrey\\Documents\\NDD\\brainlit\\docs\\notebooks\\utils\\data\n"
     ]
    }
   ],
   "source": [
    "cwd = Path(os.path.abspath(\"\"))\n",
    "data_dir = os.path.join(cwd, \"data\")\n",
    "print(f\"Downloading segments to {data_dir}\")\n",
    "if not os.path.exists(data_dir):\n",
    "    os.makedirs(data_dir)\n",
    "\n",
    "im_dir = os.path.join(data_dir, \"sample-tif-location\")\n",
    "if not os.path.exists(im_dir):\n",
    "    os.makedirs(im_dir)\n",
    "\n",
    "swc_dir = os.path.join(data_dir, \"sample-swc-location\")\n",
    "if not os.path.exists(swc_dir):\n",
    "    os.makedirs(swc_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Accessing .tif files"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "On mac/linux, we use os.path.join to construct the s3 path. However on windows you should set prefix to \"brainlit/benchmarking_data/tif-files\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "52it [00:57,  1.12s/it]\n"
     ]
    }
   ],
   "source": [
    "s3 = boto3.resource(\"s3\", config=Config(signature_version=UNSIGNED))\n",
    "bucket = s3.Bucket(\"open-neurodata\")\n",
    "prefix = \"brainlit/benchmarking_data/tif-files\"  # use this for windows\n",
    "# prefix = os.path.join(\"brainlit\", \"benchmarking_data\", \"tif-files\") #use this for mac/linux\n",
    "im_count = 0\n",
    "for _ in bucket.objects.filter(Prefix=prefix):\n",
    "    im_count += 1\n",
    "for i, im_obj in enumerate(tqdm(bucket.objects.filter(Prefix=prefix))):\n",
    "    if im_obj.key[-4:] == \".tif\":\n",
    "        im_name = os.path.basename(im_obj.key)\n",
    "        im_path = os.path.join(im_dir, im_name)\n",
    "        bucket.download_file(im_obj.key, im_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The below code can visualize a specified .tif file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "file_name = \"test_10-gfp.tif\"  # Can change to any image (test 1-25, validation 1-25)\n",
    "\n",
    "im_file = Path(im_dir) / file_name\n",
    "im = io.imread(im_file, plugin=\"tifffile\")\n",
    "\n",
    "# viewer = napari.Viewer(ndisplay=3)\n",
    "# viewer.add_image(im)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# nbscreenshot(viewer, canvas_only = True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Accessing .swc files"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Again, on windows you need to make the variable called prefix a string"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "601it [01:30,  6.67it/s]\n"
     ]
    }
   ],
   "source": [
    "s3 = boto3.resource(\"s3\", config=Config(signature_version=UNSIGNED))\n",
    "bucket = s3.Bucket(\"open-neurodata\")\n",
    "prefix = \"brainlit/benchmarking_data/Manual-GT\"  # use this for windows\n",
    "# prefix = os.path.join(\"brainlit\", \"benchmarking_data\", \"Manual-GT\") #use this for mac/linux\n",
    "swc_count = 0\n",
    "for _ in bucket.objects.filter(Prefix=prefix):\n",
    "    swc_count += 1\n",
    "for i, swc_obj in enumerate(tqdm(bucket.objects.filter(Prefix=prefix))):\n",
    "    if swc_obj.key[-4:] == \".swc\":\n",
    "        idx = swc_obj.key.find(\"Manual-GT\")\n",
    "        swc_name = swc_obj.key[idx:]\n",
    "        swc_path = os.path.join(swc_dir, swc_name)\n",
    "        dir = os.path.dirname(swc_path)\n",
    "        if not os.path.exists(dir):\n",
    "            os.makedirs(dir)\n",
    "        bucket.download_file(swc_obj.key, swc_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Aligning and visualizing images & swcs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "from brainlit.utils.benchmarking_params import (\n",
    "    brain_offsets,\n",
    "    vol_offsets,\n",
    "    scales,\n",
    "    type_to_date,\n",
    ")\n",
    "from brainlit.utils.Neuron_trace import NeuronTrace\n",
    "from pathlib import Path\n",
    "import numpy as np\n",
    "from skimage import io"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "im_dir = Path(im_dir)\n",
    "swc_base_path = Path(swc_dir) / \"Manual-GT\"\n",
    "\n",
    "gfp_files = list(im_dir.glob(\"**/*-gfp.tif\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Image 1/50\n",
      "C:\\Users\\shrey\\Documents\\NDD\\brainlit\\docs\\notebooks\\utils\\data\\sample-tif-location\\test_1-gfp.tif\n",
      "Image shape: (100, 330, 330)\n"
     ]
    }
   ],
   "source": [
    "for im_num, im_path in enumerate(gfp_files):\n",
    "    print(f\"Image {im_num+1}/{len(gfp_files)}\")\n",
    "    print(im_path)\n",
    "\n",
    "    f = im_path.parts[-1][:-8].split(\"_\")\n",
    "    image = f[0]\n",
    "    date = type_to_date[image]\n",
    "    num = int(f[1])\n",
    "\n",
    "    scale = scales[date]\n",
    "    brain_offset = brain_offsets[date]\n",
    "    vol_offset = vol_offsets[date][num]\n",
    "    im_offset = np.add(brain_offset, vol_offset)\n",
    "\n",
    "    lower = int(np.floor((num - 1) / 5) * 5 + 1)\n",
    "    upper = int(np.floor((num - 1) / 5) * 5 + 5)\n",
    "    dir1 = date + \"_\" + image + \"_\" + str(lower) + \"-\" + str(upper)\n",
    "    dir2 = date + \"_\" + image + \"_\" + str(num)\n",
    "    swc_path = swc_base_path / dir1 / dir2\n",
    "    swc_files = list(swc_path.glob(\"**/*.swc\"))\n",
    "    im = io.imread(im_path, plugin=\"tifffile\")\n",
    "    print(f\"Image shape: {im.shape}\")\n",
    "\n",
    "    paths_total = []\n",
    "    for swc_num, swc in enumerate(swc_files):\n",
    "        if \"0\" in swc.parts[-1]:\n",
    "            # skip the bounding box swc\n",
    "            continue\n",
    "\n",
    "        swc_trace = NeuronTrace(path=str(swc))\n",
    "        paths = swc_trace.get_paths()\n",
    "        swc_offset, _, _, _ = swc_trace.get_df_arguments()\n",
    "        offset_diff = np.subtract(swc_offset, im_offset)\n",
    "\n",
    "        for path_num, p in enumerate(paths):\n",
    "            pvox = (p + offset_diff) / (scale) * 1000\n",
    "            paths_total.append(pvox)\n",
    "\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# viewer = napari.Viewer(ndisplay=3)\n",
    "# viewer.add_image(np.swapaxes(im,0,2))\n",
    "# viewer.add_shapes(data=paths_total, shape_type='path', edge_width=1.0, edge_color='blue', opacity=0.8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "# nbscreenshot(viewer, canvas_only = True)"
   ]
  }
 ],
 "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}