dask_segmentation.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "tic = time.time()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/srv/conda/envs/notebook/lib/python3.7/site-packages/dipy/stats/__init__.py:7: UserWarning: The `dipy.stats` module is still under heavy development and functionality, as well as the API is likely to change in future versions of the software\n",
      "  warnings.warn(w_string)\n"
     ]
    }
   ],
   "source": [
    "import configparser\n",
    "import warnings\n",
    "import os.path as op\n",
    "import numpy as np\n",
    "import nibabel as nib\n",
    "import dipy.data as dpd\n",
    "import dipy.tracking.utils as dtu\n",
    "import dipy.tracking.streamline as dts\n",
    "from dipy.io.streamline import save_tractogram, load_tractogram\n",
    "from dipy.stats.analysis import afq_profile, gaussian_weights\n",
    "from dipy.io.stateful_tractogram import StatefulTractogram\n",
    "from dipy.io.stateful_tractogram import Space\n",
    "import dipy.core.gradients as dpg\n",
    "from dipy.reconst import dti\n",
    "from dipy.reconst import csdeconv as csd\n",
    "\n",
    "import AFQ.data as afd\n",
    "import AFQ.tractography as aft\n",
    "import AFQ.registration as reg\n",
    "import AFQ.segmentation as seg\n",
    "import AFQ.api as api\n",
    "\n",
    "import s3fs\n",
    "\n",
    "import logging\n",
    "logging.basicConfig(level=logging.INFO)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "CP = configparser.ConfigParser()\n",
    "CP.read_file(open(op.join(op.expanduser('~'), '.aws', 'credentials')))\n",
    "hcp_ak = CP.get('hcp', 'AWS_ACCESS_KEY_ID')\n",
    "hcp_sk = CP.get('hcp', 'AWS_SECRET_ACCESS_KEY')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "fs = s3fs.S3FileSystem(key=hcp_ak, secret=hcp_sk)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "CP = configparser.ConfigParser()\n",
    "CP.read_file(open(op.join(op.expanduser('~'), '.aws', 'credentials')))\n",
    "ak = CP.get('default', 'AWS_ACCESS_KEY_ID')\n",
    "sk = CP.get('default', 'AWS_SECRET_ACCESS_KEY')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "my_fs =  s3fs.S3FileSystem(key=ak, secret=sk)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "subject = 100307"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# dwi_fname = f'hcp-openaccess/HCP_1200/{subject}/T1w/Diffusion/data.nii.gz'\n",
    "# dwi_img = afd.s3fs_nifti_read(dwi_fname, fs=fs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# with fs.open(f'hcp-openaccess/HCP_1200/{subject}/T1w/Diffusion/bvals') as ff:\n",
    "#     bvals = np.loadtxt(ff)\n",
    "    \n",
    "# with fs.open(f'hcp-openaccess/HCP_1200/{subject}/T1w/Diffusion/bvecs') as ff:\n",
    "#     bvecs = np.loadtxt(ff)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# gtab = dpg.gradient_table(bvals, bvecs, b0_threshold=50)\n",
    "# mapping = reg.syn_register_dwi(dwi_img, gtab)[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# reg.write_mapping(mapping, 'mapping.nii.gz')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# mapping_img = nib.load('mapping.nii.gz')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "# afd.s3fs_nifti_write(mapping_img, \n",
    "#                      f'hcp.pangeo.experiments/{subject}/mapping.nii.gz', \n",
    "#                      fs=my_fs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "bundle_names = api.BUNDLES"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def segment_bundle(sl_idx):\n",
    "    log = logging.getLogger(__name__)\n",
    "    with fs.open(f'hcp-openaccess/HCP_1200/{subject}/T1w/Diffusion/bvals') as ff:\n",
    "        bvals = np.loadtxt(ff)\n",
    "        np.savetxt('bvals', bvals)\n",
    "\n",
    "    with fs.open(f'hcp-openaccess/HCP_1200/{subject}/T1w/Diffusion/bvecs') as ff:\n",
    "        bvecs = np.loadtxt(ff)\n",
    "        np.savetxt('bvecs', bvecs)\n",
    "\n",
    "    log.info(\"Getting DWI data\")\n",
    "    dwi_fname = f'hcp-openaccess/HCP_1200/{subject}/T1w/Diffusion/data.nii.gz'\n",
    "    dwi_img = afd.s3fs_nifti_read(dwi_fname, fs=fs)\n",
    "    \n",
    "    log.info(\"Getting mapping\")\n",
    "    mapping_fname = f'hcp.pangeo.experiments/{subject}/mapping.nii.gz'\n",
    "    mapping_img = afd.s3fs_nifti_read(mapping_fname, fs=my_fs)\n",
    "    reg_template =  dpd.read_mni_template()\n",
    "    mapping = reg.read_mapping(mapping_img, dwi_img, reg_template)\n",
    "    \n",
    "    log.info(\"Getting bundle dict\")\n",
    "    bundle_dict = api.make_bundle_dict(bundle_names=bundle_names)\n",
    "    AFQ = seg.Segmentation()\n",
    "\n",
    "    sl_file = f'hcp.pangeo.experiments/{subject}/sl-{sl_idx:03d}.trk'\n",
    "    log.info(f\"Segmenting {sl_file}\")\n",
    "    my_fs.download(sl_file, 'tmp.trk')\n",
    "    tg = load_tractogram('./tmp.trk', dwi_img)\n",
    "\n",
    "    streamlines = dts.Streamlines(\n",
    "        dtu.transform_tracking_output(tg.streamlines,\n",
    "                                  np.linalg.inv(dwi_img.affine)))\n",
    "\n",
    "    fiber_groups = AFQ.segment(bundle_dict, \n",
    "                               streamlines, \n",
    "                               fdata=dwi_img, \n",
    "                               fbval='bvals', \n",
    "                               fbvec='bvecs',\n",
    "                               mapping=mapping,\n",
    "                               b0_threshold=50)\n",
    "\n",
    "    log.info(\"Saving out results and uploading\")\n",
    "    for kk in fiber_groups:\n",
    "        if len(fiber_groups[kk]) > 0:\n",
    "            print(kk, len(fiber_groups[kk]))\n",
    "\n",
    "            sft = StatefulTractogram(\n",
    "                dtu.transform_tracking_output(fiber_groups[kk], dwi_img.affine),\n",
    "                dwi_img, \n",
    "                Space.RASMM)\n",
    "\n",
    "            save_tractogram(sft, f'./{kk}-{sl_idx}_afq.trk',\n",
    "                            bbox_valid_check=False)\n",
    "\n",
    "            my_fs.upload(f'./{kk}-{sl_idx}_afq.trk', \n",
    "                         f'hcp.pangeo.experiments/{subject}/{kk}-c_afq.trk')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:__main__:Getting DWI data\n",
      "INFO:__main__:Getting mapping\n",
      "pixdim[0] (qfac) should be 1 (default) or -1; setting qfac to 1\n",
      "INFO:nibabel.global:pixdim[0] (qfac) should be 1 (default) or -1; setting qfac to 1\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Data size is approximately 70MB\n",
      "Dataset is already in place. If you want to fetch it again please first remove the folder /home/jovyan/.dipy/mni_template \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:__main__:Getting bundle dict\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset is already in place. If you want to fetch it again please first remove the folder /home/jovyan/AFQ_data/templates \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:__main__:Segmenting hcp.pangeo.experiments/100307/sl-000.trk\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset is already in place. If you want to fetch it again please first remove the folder /home/jovyan/AFQ_data/callosum_templates \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:Moved streamlines from rasmm to vox\n",
      "INFO:root:Origin moved to the center of voxel\n",
      "INFO:root:Moved streamlines from vox to rasmm\n",
      "INFO:root:Origin moved to the corner of voxel\n",
      "INFO:AFQ.Segmentation:Preparing image, bvals, bvecs...\n",
      "INFO:AFQ.Segmentation:Preparing mapping...\n",
      "pixdim[0] (qfac) should be 1 (default) or -1; setting qfac to 1\n",
      "INFO:nibabel.global:pixdim[0] (qfac) should be 1 (default) or -1; setting qfac to 1\n",
      "INFO:AFQ.Segmentation:Preprocessing Streamlines...\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Data size is approximately 70MB\n",
      "Dataset is already in place. If you want to fetch it again please first remove the folder /home/jovyan/.dipy/mni_template \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:AFQ.Segmentation:Assigning Streamlines to Bundles...\n",
      "INFO:AFQ.Segmentation:Bundle: ATR_R\n",
      "100%|██████████| 32592/32592 [00:17<00:00, 1835.49it/s]\n",
      "INFO:AFQ.Segmentation:Bundle: ATR_L\n",
      "100%|██████████| 32592/32592 [00:00<00:00, 283017.89it/s]\n",
      "INFO:AFQ.Segmentation:Bundle: CGC_R\n",
      "100%|██████████| 32592/32592 [00:09<00:00, 3392.75it/s]\n",
      "INFO:AFQ.Segmentation:Bundle: CGC_L\n",
      "100%|██████████| 32592/32592 [00:00<00:00, 536241.22it/s]\n",
      "INFO:AFQ.Segmentation:Bundle: CST_R\n",
      "100%|██████████| 32592/32592 [00:03<00:00, 10078.89it/s]\n",
      "INFO:AFQ.Segmentation:Bundle: CST_L\n",
      "100%|██████████| 32592/32592 [00:01<00:00, 22282.82it/s]\n",
      "INFO:AFQ.Segmentation:Bundle: HCC_R\n",
      "100%|██████████| 32592/32592 [00:01<00:00, 18266.30it/s]\n",
      "INFO:AFQ.Segmentation:Bundle: HCC_L\n",
      "100%|██████████| 32592/32592 [00:00<00:00, 507385.28it/s]\n",
      "INFO:AFQ.Segmentation:Bundle: IFO_R\n"
     ]
    }
   ],
   "source": [
    "segment_bundle(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "isinstance()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dask.distributed import Client\n",
    "from dask_kubernetes import KubeCluster"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_workers = 40"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cluster = KubeCluster(n_workers=n_workers)\n",
    "cluster"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "client = Client(cluster)\n",
    "client"
   ]
  },
  {
   "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}