{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This demo presents the usage of NiftyNet's image sampler implemented at `niftynet.contrib.dataset_sampler.*`\n",
    "\n",
    "What's a `sampler`?\n",
    "\n",
    "The sampler takes an image reader as input, produces image windows sampled from each image.\n",
    "\n",
    "Why `sampler`?\n",
    "\n",
    "In many cases in medical imaging applications, deep neural networks operate on image windows instead of the entire images/datasets, partly due to the GPU memory limitation and training efficiencies. The sampler provides ready-to-use interfaces to generate image windows, works well with `tf.data.Dataset`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Before the demo..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First make sure the source code is available, and import the module.\n",
    "\n",
    "For NiftyNet installation, please checkout:\n",
    "\n",
    "http://niftynet.readthedocs.io/en/dev/installation.html\n",
    "\n",
    "For demonstration purpose we download some demo data to `~/niftynet/data/`\n",
    "\n",
    "For visualisation we install `matplotlib`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 158,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accessing: https://cmiclab.cs.ucl.ac.uk/CMIC/NiftyNetExampleServer\n",
      "mr_ct_regression_model_zoo_data: OK. \n",
      "Already downloaded. Use the -r option to download again.\n",
      "Requirement already satisfied: matplotlib in /usr/local/anaconda3/envs/tf/lib/python3.6/site-packages (2.2.2)\n",
      "Requirement already satisfied: pytz in /usr/local/anaconda3/envs/tf/lib/python3.6/site-packages (from matplotlib) (2018.4)\n",
      "Requirement already satisfied: cycler>=0.10 in /usr/local/anaconda3/envs/tf/lib/python3.6/site-packages (from matplotlib) (0.10.0)\n",
      "Requirement already satisfied: python-dateutil>=2.1 in /usr/local/anaconda3/envs/tf/lib/python3.6/site-packages (from matplotlib) (2.7.2)\n",
      "Requirement already satisfied: numpy>=1.7.1 in /usr/local/anaconda3/envs/tf/lib/python3.6/site-packages (from matplotlib) (1.14.3)\n",
      "Requirement already satisfied: kiwisolver>=1.0.1 in /usr/local/anaconda3/envs/tf/lib/python3.6/site-packages (from matplotlib) (1.0.1)\n",
      "Requirement already satisfied: pyparsing!=2.0.4,!=2.1.2,!=2.1.6,>=2.0.1 in /usr/local/anaconda3/envs/tf/lib/python3.6/site-packages (from matplotlib) (2.2.0)\n",
      "Requirement already satisfied: six>=1.10 in /usr/local/anaconda3/envs/tf/lib/python3.6/site-packages (from matplotlib) (1.11.0)\n",
      "Requirement already satisfied: setuptools in /usr/local/anaconda3/envs/tf/lib/python3.6/site-packages (from kiwisolver>=1.0.1->matplotlib) (39.1.0)\n"
     ]
    }
   ],
   "source": [
    "import sys\n",
    "niftynet_path = '/Users/foo/Documents/Niftynet/'\n",
    "sys.path.append(niftynet_path)\n",
    "\n",
    "from niftynet.utilities.download import download\n",
    "download('mr_ct_regression_model_zoo_data')\n",
    "\n",
    "!{sys.executable} -m pip install matplotlib"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Image as 'window'\n",
    "\n",
    "The simplest use case is treating each image as an image window. This is implemented by default using `ImageWindowDataset` class from `niftynet.contrib.dataset_sampler`. This class also acts as a base class, which can be extended to generate smaller windows using different sampling strategies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 159,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1mINFO:niftynet:\u001b[0m \n",
      "\n",
      "Number of subjects 15, input section names: ['subject_id', 'CT']\n",
      "-- using all subjects (without data partitioning).\n",
      "\n",
      "\u001b[1mINFO:niftynet:\u001b[0m Image reader: loading 15 subjects from sections ('CT',) as input [CT]\n"
     ]
    }
   ],
   "source": [
    "from niftynet.io.image_reader import ImageReader\n",
    "from niftynet.contrib.dataset_sampler.image_window_dataset import ImageWindowDataset\n",
    "\n",
    "# creating an image reader.\n",
    "data_param = \\\n",
    "    {'CT': {'path_to_search': '~/niftynet/data/mr_ct_regression/CT_zero_mean',\n",
    "            'filename_contains': 'nii'}}\n",
    "reader = ImageReader().initialise(data_param)\n",
    "\n",
    "# creating a window sampler dataset from the reader\n",
    "sampler = ImageWindowDataset(reader)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The sampler can be used as a numpy function, or a tensorflow operation.\n",
    "### Use the sampler as a numpy function\n",
    "Directly call the instance (this is actually invoking `sampler.layer_op`):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 160,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dict_keys(['CT', 'CT_location']) [[  0   0   0   0 288 288   1]] (1, 288, 288, 1, 1, 1)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 40884 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "windows = sampler()\n",
    "print(windows.keys(), windows['CT_location'], windows['CT'].shape)\n",
    "    \n",
    "import matplotlib.pyplot as plt\n",
    "plt.imshow(windows['CT'][0,:,:,0,0,0])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Use the sampler as a tensorflow op\n",
    "First add a iterator node, this is wrapped as `pop_batch_op()`,\n",
    "then run the op."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 161,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dict_keys(['CT', 'CT_location']) [[  5   0   0   0 288 288   1]] (1, 288, 288, 1, 1)\n"
     ]
    }
   ],
   "source": [
    "import tensorflow as tf\n",
    "# adding the tensorflow tensors\n",
    "next_window = sampler.pop_batch_op()\n",
    "\n",
    "# run the tensors\n",
    "with tf.Session() as sess:\n",
    "    sampler.run_threads(sess) #initialise the iterator\n",
    "    windows = sess.run(next_window)\n",
    "    print(windows.keys(), windows['CT_location'], windows['CT'].shape)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- The location array `['MR_location']` represents the spatial coordinates of the window: \n",
    "  \n",
    "  `[subject_id, x_start, y_start, z_start, x_end, y_end, z_end]`\n",
    "\n",
    "\n",
    "- As a numpy function, the output shape is `(1, x, y, z, 1, 1)` which represents batch, width, height, depth, time points, channels. As a tensorflow op, the output shape is `(batch, x,[ y,][ z,] channels)`, which means the time dimension (not currently supported) and the spatial axes will be \"squeezed\" if the length along them is one. This simplifies the network definitions based on 2D or 3D outputs of the sampler.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Uniform sampler\n",
    "Generating image windows randomly from images.\n",
    "\n",
    "This is implemented by overriding the `layer_op` of `ImageWindowDataset`.\n",
    "\n",
    "The following code creates a uniform sampler and draws an image window.\n",
    "\n",
    "Two sections `MR` and `CT` are given as the input data parameter,\n",
    "the reader loads these sections by matching the filenames, and\n",
    "outputs the sampling windows (from the same spatial coordinates for all input images) as a dictionary.\n",
    "\n",
    "- When the spatial window sizes are different in `MR` and `CT`, concentric windows will be sampled.\n",
    "\n",
    "- When the spatial window size is (0, 0, 0), size of the first image from the reader will be used as the window size.\n",
    "\n",
    "- `windows_per_image` parameter specifies the number of image windows from each image returned by the reader."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 169,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1mINFO:niftynet:\u001b[0m \n",
      "\n",
      "Number of subjects 15, input section names: ['subject_id', 'MR', 'CT']\n",
      "-- using all subjects (without data partitioning).\n",
      "\n",
      "\u001b[1mINFO:niftynet:\u001b[0m Image reader: loading 15 subjects from sections ('MR',) as input [MR]\n",
      "\u001b[1mINFO:niftynet:\u001b[0m Image reader: loading 15 subjects from sections ('CT',) as input [CT]\n",
      "\u001b[1mINFO:niftynet:\u001b[0m initialised uniform sampler {'MR': (1, 80, 80, 1, 1, 1), 'MR_location': (1, 7), 'CT': (1, 80, 80, 1, 1, 1), 'CT_location': (1, 7)} \n",
      "[[  2  38  66   0 118 146   1]\n",
      " [  1  90 164   0 170 244   1]] (2, 80, 80, 1)\n",
      "[[  2  38  66   0 118 146   1]\n",
      " [  1  90 164   0 170 244   1]] (2, 80, 80, 1)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 41572 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from niftynet.io.image_reader import ImageReader\n",
    "from niftynet.contrib.dataset_sampler.sampler_uniform_v2 import UniformSampler\n",
    "\n",
    "# creating an image reader.\n",
    "# creating an image reader.\n",
    "data_param = \\\n",
    "    {'MR': {'path_to_search': '~/niftynet/data/mr_ct_regression/CT_zero_mean',\n",
    "            'filename_contains': 'nii',\n",
    "            'spatial_window_size': (80, 80, 1)},\n",
    "     'CT': {'path_to_search': '~/niftynet/data/mr_ct_regression/T2_corrected',\n",
    "            'filename_contains': 'nii',\n",
    "            'spatial_window_size': (80, 80, 1)},\n",
    "    }\n",
    "window_sizes = {'MR': (80, 80, 1), 'CT': (80, 80, 1)}\n",
    "reader = ImageReader().initialise(data_param)\n",
    "\n",
    "# uniform sampler returns windows of 32^3-voxels\n",
    "uniform_sampler = UniformSampler(\n",
    "    reader, window_sizes, batch_size=2, windows_per_image=5)\n",
    "\n",
    "\n",
    "import tensorflow as tf\n",
    "# adding the tensorflow tensors\n",
    "next_window = uniform_sampler.pop_batch_op()\n",
    "\n",
    "# run the tensors\n",
    "with tf.Session() as sess:\n",
    "    uniform_sampler.run_threads(sess) #initialise the iterator\n",
    "    windows = sess.run(next_window)\n",
    "    print(windows['MR_location'], windows['MR'].shape)\n",
    "    print(windows['CT_location'], windows['CT'].shape)\n",
    "\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "plt.figure()\n",
    "plt.subplot(1,2,1)\n",
    "plt.imshow(windows['MR'][0,:,:,0])\n",
    "plt.subplot(1,2,2)\n",
    "plt.imshow(windows['CT'][0,:,:,0])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Grid sampler\n",
    "Generatoring image windows from images with a sliding window.\n",
    "\n",
    "This is implemented by overriding the `layer_op` of `ImageWindowDataset`.\n",
    "\n",
    "The `window_border` parameter controls the amount of overlap in between sampling windows.\n",
    "When the grid sampler is used for fully convolutional inference, the overlapping regions of the windows are cropped, so that the windows can be aggregated as an image at the resolution of the input.  \n",
    "\n",
    "See also: http://niftynet.readthedocs.io/en/dev/config_spec.html#border"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 163,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1mINFO:niftynet:\u001b[0m \n",
      "\n",
      "Number of subjects 15, input section names: ['subject_id', 'CT']\n",
      "-- using all subjects (without data partitioning).\n",
      "\n",
      "\u001b[1mINFO:niftynet:\u001b[0m Image reader: loading 15 subjects from sections ('CT',) as input [CT]\n",
      "\u001b[1mINFO:niftynet:\u001b[0m initialised window instance\n",
      "\u001b[1mINFO:niftynet:\u001b[0m initialised grid sampler {'CT': (1, 42, 42, 1, 1, 1), 'CT_location': (1, 7)}\n",
      "\u001b[1mINFO:niftynet:\u001b[0m grid sampling image sizes: {'CT': (288, 288, 1, 1, 1)}\n",
      "\u001b[1mINFO:niftynet:\u001b[0m grid sampling window sizes: {'CT': (42, 42, 1, 1, 1)}\n",
      "\u001b[1mINFO:niftynet:\u001b[0m yielding 121 locations from image\n",
      "\u001b[1mINFO:niftynet:\u001b[0m grid sampling image sizes: {'CT': (288, 288, 1, 1, 1)}\n",
      "\u001b[1mINFO:niftynet:\u001b[0m grid sampling window sizes: {'CT': (42, 42, 1, 1, 1)}\n",
      "\u001b[1mINFO:niftynet:\u001b[0m yielding 121 locations from image\n"
     ]
    }
   ],
   "source": [
    "from niftynet.io.image_reader import ImageReader\n",
    "from niftynet.contrib.dataset_sampler.sampler_grid_v2 import GridSampler\n",
    "\n",
    "# creating an image reader.\n",
    "data_param = \\\n",
    "    {'CT': {'path_to_search': '~/niftynet/data/mr_ct_regression/CT_zero_mean',\n",
    "            'filename_contains': 'nii'}}\n",
    "reader = ImageReader().initialise(data_param)\n",
    "\n",
    "# uniform sampler returns windows of 32^3-voxels\n",
    "uniform_sampler = GridSampler(reader, \n",
    "                              window_sizes=(42, 42, 1), \n",
    "                              window_border=(8,8,1), batch_size=1)\n",
    "\n",
    "\n",
    "import tensorflow as tf\n",
    "# adding the tensorflow tensors\n",
    "next_window = uniform_sampler.pop_batch_op()\n",
    "\n",
    "# run the tensors\n",
    "with tf.Session() as sess:\n",
    "    uniform_sampler.run_threads(sess) #initialise the iterator\n",
    "    subject_id = 0\n",
    "    coords = []\n",
    "    while True:\n",
    "        windows = sess.run(next_window)\n",
    "        if not subject_id == windows['CT_location'][0,0]:\n",
    "            break;\n",
    "        #print(windows.keys(), windows['MR_location'], windows['MR'].shape)\n",
    "        coords.append(windows['CT_location'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Visualisation of the window coordinates (change `window_sizes` and `window_border` to see different window allocations):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 164,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(288, 288, 1, 1, 1)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 29640 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.patches as patches\n",
    "from matplotlib.collections import PatchCollection\n",
    "\n",
    "f, (ax1, ax) = plt.subplots(1,2)\n",
    "\n",
    "# show image\n",
    "_, img, _ = reader(idx=0)\n",
    "print(img['CT'].shape)\n",
    "plt.subplot(1,2,1)\n",
    "plt.imshow(img['CT'][:,:,0,0,0])\n",
    "\n",
    "# show sampled windows\n",
    "all_patch = []\n",
    "for win in np.concatenate(coords, axis=0):\n",
    "    patch = patches.Rectangle(\n",
    "        (win[1], win[2]),\n",
    "        win[4]-win[1], win[5]-win[2], linewidth=1)\n",
    "    all_patch.append(patch)\n",
    "all_pc = PatchCollection(\n",
    "    all_patch, alpha=0.1, edgecolor='r', facecolor='r')\n",
    "ax.add_collection(all_pc)\n",
    "ax.set_xlim([0, np.max(coords, axis=0)[0,4]])\n",
    "ax.set_ylim([0, np.max(coords, axis=0)[0,5]])\n",
    "ax.set_aspect('equal', 'datalim')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Weighted sampler\n",
    "Generatoring image windows from images with a sampling prior of the foreground.\n",
    "\n",
    "This sampler uses a cumulative histogram for fast sampling, works with both continous and discrete maps.\n",
    "It is implemented by overriding the `layer_op` of `ImageWindowDataset`.\n",
    "\n",
    "Weight map can be specified by an input specification with `sampler` as the key.\n",
    "\n",
    "The following example uses a binary mask as the weight map."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 165,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1mINFO:niftynet:\u001b[0m \n",
      "\n",
      "Number of subjects 15, input section names: ['subject_id', 'CT', 'sampler']\n",
      "-- using all subjects (without data partitioning).\n",
      "\n",
      "\u001b[1mINFO:niftynet:\u001b[0m Image reader: loading 1 subjects from sections ('CT',) as input [CT]\n",
      "\u001b[1mINFO:niftynet:\u001b[0m Image reader: loading 1 subjects from sections ('sampler',) as input [sampler]\n",
      "\u001b[1mINFO:niftynet:\u001b[0m initialised uniform sampler {'CT': (100, 12, 12, 1, 1, 1), 'CT_location': (100, 7), 'sampler': (100, 12, 12, 1, 1, 1), 'sampler_location': (100, 7)} \n",
      "\u001b[1mINFO:niftynet:\u001b[0m Initialised weighted sampler window instance\n"
     ]
    }
   ],
   "source": [
    "from niftynet.io.image_reader import ImageReader\n",
    "from niftynet.contrib.dataset_sampler.sampler_weighted_v2 import WeightedSampler\n",
    "\n",
    "# creating an image reader.\n",
    "data_param = \\\n",
    "    {'CT':      {'path_to_search': '~/niftynet/data/mr_ct_regression/CT_zero_mean',\n",
    "                 'filename_contains': 'PAR.nii.gz'},\n",
    "     'sampler': {'path_to_search': '~/niftynet/data/mr_ct_regression/T2_mask',\n",
    "                 'filename_contains': 'nii'}}\n",
    "reader = ImageReader().initialise(data_param)\n",
    "\n",
    "weighted_sampler = WeightedSampler(\n",
    "    reader, window_sizes=(12, 12, 1), batch_size=1, windows_per_image=100)\n",
    "\n",
    "import tensorflow as tf\n",
    "# adding the tensorflow tensors\n",
    "next_window = weighted_sampler.pop_batch_op()\n",
    "\n",
    "# run the tensors\n",
    "with tf.Session() as sess:\n",
    "    weighted_sampler.run_threads(sess) #initialise the iterator\n",
    "    coords = []\n",
    "    for _ in range(200):\n",
    "        windows = sess.run(next_window)\n",
    "        #print(windows.keys(), windows['CT_location'], windows['CT'].shape)\n",
    "        coords.append(windows['CT_location'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 166,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(288, 288, 1, 1, 1)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 34024 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.patches as patches\n",
    "from matplotlib.collections import PatchCollection\n",
    "\n",
    "f, (ax1, ax) = plt.subplots(1,2)\n",
    "\n",
    "# show image\n",
    "_, img, _ = reader(idx=0)\n",
    "print(img['CT'].shape)\n",
    "plt.subplot(1,2,1)\n",
    "plt.imshow(img['CT'][:,:,0,0,0])\n",
    "#plt.subplot(1,2,2)\n",
    "ax.imshow(img['sampler'][:,:,0,0,0])\n",
    "\n",
    "# show sampled windows\n",
    "all_patch = []\n",
    "for win in np.concatenate(coords, axis=0):\n",
    "    patch = patches.Rectangle(\n",
    "        (win[2], win[1]),\n",
    "        win[5]-win[2], win[4]-win[1], linewidth=1)\n",
    "    all_patch.append(patch)\n",
    "all_pc = PatchCollection(\n",
    "    all_patch, alpha=0.5, edgecolor='r', facecolor='r')\n",
    "ax.add_collection(all_pc)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Balanced sampler\n",
    "Generatoring image windows from images with a sampling prior of the foreground.\n",
    "\n",
    "This sampler generates image windows from a discrete label map as if every label\n",
    "had the same probability of occurrence.\n",
    "It is implemented by overriding the `layer_op` of `ImageWindowDataset`.\n",
    "\n",
    "Weight map can be specified by an input specification with `sampler` as the key.\n",
    "The following example uses a foreground mask as the weight map."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 167,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1mINFO:niftynet:\u001b[0m \n",
      "\n",
      "Number of subjects 15, input section names: ['subject_id', 'CT', 'sampler']\n",
      "-- using all subjects (without data partitioning).\n",
      "\n",
      "\u001b[1mINFO:niftynet:\u001b[0m Image reader: loading 1 subjects from sections ('CT',) as input [CT]\n",
      "\u001b[1mINFO:niftynet:\u001b[0m Image reader: loading 1 subjects from sections ('sampler',) as input [sampler]\n",
      "\u001b[1mINFO:niftynet:\u001b[0m initialised uniform sampler {'CT': (100, 12, 12, 1, 1, 1), 'CT_location': (100, 7), 'sampler': (100, 12, 12, 1, 1, 1), 'sampler_location': (100, 7)} \n",
      "\u001b[1mINFO:niftynet:\u001b[0m Initialised balanced sampler window instance\n"
     ]
    }
   ],
   "source": [
    "from niftynet.io.image_reader import ImageReader\n",
    "from niftynet.contrib.dataset_sampler.sampler_balanced_v2 import BalancedSampler\n",
    "\n",
    "# creating an image reader.\n",
    "data_param = \\\n",
    "    {'CT':      {'path_to_search': '~/niftynet/data/mr_ct_regression/CT_zero_mean',\n",
    "                 'filename_contains': 'PAR.nii.gz'},\n",
    "     'sampler': {'path_to_search': '~/niftynet/data/mr_ct_regression/T2_mask',\n",
    "                 'filename_contains': 'nii'}}\n",
    "reader = ImageReader().initialise(data_param)\n",
    "\n",
    "balanced_sampler = BalancedSampler(\n",
    "    reader, window_sizes=(12, 12, 1), windows_per_image=100)\n",
    "\n",
    "import tensorflow as tf\n",
    "# adding the tensorflow tensors\n",
    "next_window = balanced_sampler.pop_batch_op()\n",
    "\n",
    "# run the tensors\n",
    "with tf.Session() as sess:\n",
    "    balanced_sampler.run_threads(sess) #initialise the iterator\n",
    "    coords = []\n",
    "    for _ in range(200):\n",
    "        windows = sess.run(next_window)\n",
    "        #print(windows.keys(), windows['CT_location'], windows['CT'].shape)\n",
    "        coords.append(windows['CT_location'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Visualisation of the window coordinates (change `data_param` see different window allocations):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 151,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(288, 288, 1, 1, 1)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 33240 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.patches as patches\n",
    "from matplotlib.collections import PatchCollection\n",
    "\n",
    "f, (ax1, ax) = plt.subplots(1,2)\n",
    "\n",
    "# show image\n",
    "_, img, _ = reader(idx=0)\n",
    "print(img['CT'].shape)\n",
    "plt.subplot(1,2,1)\n",
    "plt.imshow(img['CT'][:,:,0,0,0])\n",
    "#plt.subplot(1,2,2)\n",
    "ax.imshow(img['sampler'][:,:,0,0,0])\n",
    "\n",
    "# show sampled windows\n",
    "all_patch = []\n",
    "for win in np.concatenate(coords, axis=0):\n",
    "    patch = patches.Rectangle(\n",
    "        (win[2], win[1]),\n",
    "        win[5]-win[2], win[4]-win[1], linewidth=1)\n",
    "    all_patch.append(patch)\n",
    "all_pc = PatchCollection(\n",
    "    all_patch, alpha=0.5, edgecolor='r', facecolor='r')\n",
    "ax.add_collection(all_pc)\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:tf]",
   "language": "python",
   "name": "conda-env-tf-py"
  },
  "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.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
