{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Cell Segmentation Using Spherical Harmonics\n",
    "\n",
    "\n",
    "The image data used for this project is published in https://www.pnas.org/content/113/51/E8238 with the corresponding repository at https://www.repository.cam.ac.uk/handle/1810/262530.\n",
    "\n",
    "\n",
    "## Instance Encoding\n",
    "To demonstrate how cell segmentations can be encoded by spherical harmonic (SH) coefficients, an examplary cropped patch from the public data set is used."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "from skimage import io\n",
    "from dipy.core.geometry import sphere2cart, cart2sphere\n",
    "from keras.models import model_from_json\n",
    "\n",
    "from util.utils import get_sampling_sphere, spherical_instance_sampling, sampling2harmonics, harmonics2sampling, \\\n",
    "                       agglomerative_clustering, harmonic_non_max_suppression\n",
    "from util.data_handling import descriptors2image, instances2harmonicmask, harmonicmask2sampling\n",
    "from models.networks import harmonic_rcnn_3D"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the image and the corresponding mask\n",
    "img = io.imread('demo_img.tif')\n",
    "seg = io.imread('demo_seg.tif')\n",
    "\n",
    "# Reorder to have (x,y,z)\n",
    "img = np.transpose(img, (2,1,0))\n",
    "seg = np.transpose(seg, (2,1,0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 28016 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Plot the image and mask\n",
    "plt.subplot(1,2,1)\n",
    "plt.imshow(img[...,32])\n",
    "plt.subplot(1,2,2)\n",
    "plt.imshow(seg[...,32])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To sample each instance and encode it to a compact representation, sample positions have to be specified. Those positions are chosen by following the principle of electrostatic repulsion and are determined in an iterative manner. For the proposed method, 5000 sample positions are optimized within 10000 iterations. For demonstration purposes, the acquisition of 200 angular positions within only 1000 iterations is shown in the following."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 90296 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "sampling_sphere = get_sampling_sphere(num_sample_points=200, num_iterations=1000, plot_sampling=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The final sampling pattern is derived from the $\\theta$ and $\\phi$ values of each spherical coordinate.\n",
    "Since the calculation of a more precice sampling grid would take too long, a pre-calculated set of 5000 sample positions is loaded."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load angular pattern\n",
    "theta_phi_sampling = np.load('theta_phi_sampling_5000points_10000iter.npy')\n",
    "\n",
    "# Sample each cell instance in the image crop\n",
    "instances, sampled_radii, centroids = spherical_instance_sampling(seg, theta_phi_sampling, bg_values=[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For each of the instances, information about the instance label, the length of the sampled radius at each angular position and the location of the centroid were derived. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The first cell has label 516 and its centroid is located at [57.40485437 29.48252427 58.62135922].\n",
      "The sampled radii are a vector of length 5000, which corresponds to the total number of angular directions.\n"
     ]
    }
   ],
   "source": [
    "print('The first cell has label {0} and its centroid is located at {1}.\\n'.format(instances[0], centroids[0])+\\\n",
    "      'The sampled radii are a vector of length {0}, which corresponds to '.format(len(sampled_radii[0]))+\\\n",
    "      'the total number of angular directions.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Based on this sampling, spherical harmonic coefficients can be computed. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up the converter\n",
    "# First input is the SH order (order 8 = 81 coefficients)\n",
    "s2h_converter = sampling2harmonics(8, theta_phi_sampling) \n",
    "\n",
    "# Encode each instance segmentation \n",
    "harmonic_sampling = s2h_converter.convert(sampled_radii)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ 2.75625377e+01,  7.36752057e-01, -1.70148074e+00, -1.63059714e+00,\n",
       "       -1.77123035e+00,  1.46504441e+00, -2.93435339e+00,  2.29121719e+00,\n",
       "        3.57302708e-02,  2.63368373e+00, -2.16214198e-01, -1.51947047e+00,\n",
       "       -2.45226467e+00, -1.10525302e+00,  4.51669330e-01,  8.25379156e-01,\n",
       "        2.77154744e-01, -4.33032983e-01,  6.13681547e-01, -3.33502420e-01,\n",
       "        1.64150957e-01, -3.31935032e-01, -3.38347478e-01, -4.70662169e-01,\n",
       "        3.95907652e-02, -6.18516742e-01,  1.58342622e-01, -6.11273779e-01,\n",
       "        1.28662602e-01,  1.68269311e-01,  1.33994741e+00,  1.56610975e-01,\n",
       "       -4.01387128e-02, -2.72221286e-02,  4.30783268e-01, -2.01465358e-01,\n",
       "        4.45510409e-01, -1.17482156e-01, -3.72896887e-01,  4.41898210e-02,\n",
       "        4.26253409e-01,  4.01642346e-01,  4.94688806e-01,  3.59734783e-01,\n",
       "       -8.57002796e-02,  3.05929087e-02, -3.87996983e-01,  6.44160422e-02,\n",
       "        3.24784649e-02,  2.46687356e-02, -9.27466632e-02,  2.92219664e-01,\n",
       "        4.24961328e-02, -9.41522883e-02, -4.89565598e-01,  1.81967364e-01,\n",
       "       -2.67550424e-01, -2.55166444e-02,  1.98729762e-01,  1.21653575e-02,\n",
       "       -4.75068479e-02, -2.89425764e-03, -3.25157234e-01,  5.95464073e-03,\n",
       "       -1.63164883e-01,  3.56173444e-02, -9.04498200e-02,  2.76659142e-02,\n",
       "        3.27613774e-02,  2.27555427e-01, -2.18438739e-01, -1.72686282e-01,\n",
       "       -4.44378007e-01, -1.36038190e-01,  1.84336120e-02,  5.37255257e-02,\n",
       "        1.23297448e-01, -5.37147444e-02,  6.52444648e-02,  7.54067746e-02,\n",
       "       -9.74218939e-02])"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "harmonic_sampling[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Instance Decoding\n",
    "To reverse the encoding and obtain voxelized segmentation masks, the harmonic representation has to be decoded using the radius sampling pattern, which ultimatly serves as input for a Delaunay triangulation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up the converter for decoding\n",
    "h2s_converter = harmonics2sampling(8, theta_phi_sampling)\n",
    "\n",
    "# Decode the instances\n",
    "decoded_radii = h2s_converter.convert(harmonic_sampling)\n",
    "\n",
    "# Assume certain predictions\n",
    "prob_objects = [1,]*len(decoded_radii) \n",
    "\n",
    "# Reconstruct the mask via Delaunay triangulation\n",
    "decoded_mask = descriptors2image([centroids, prob_objects, decoded_radii], theta_phi_sampling=theta_phi_sampling, shape=(64,64,64))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 28576 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Plot the image and mask\n",
    "plt.subplot(1,2,1)\n",
    "plt.imshow(img[...,32])\n",
    "plt.subplot(1,2,2)\n",
    "plt.imshow(decoded_mask[...,32])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Instance Segmentation Using SH\n",
    "\n",
    "In order to use a CNN for prediction of cell shapes, a mask is constructed, which illustratively divides the input image into different segments and encodes the position, confidence and shape of possible cell instances in each of those segments. Considering the image patch, the resulting mask will be constructed in the following."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(8, 8, 8, 170)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "encoded_mask = instances2harmonicmask(seg[...,np.newaxis], s2h_converter, shape=(64,64,64), cell_size=(8,8,8), dets_per_region=2, bg_values=[1])\n",
    "encoded_mask.shape\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The segment size is $8\\times8\\times8$ pixel, which is constrained by the 3 conv layers with stride=2 in the network architecture. Since a patch size of $64\\times64\\times64$ pixel was used, the mask has a spatial size of  $8\\times8\\times8$ pixel. Due to the employment of harmonic order 8, each shape is encoded using a total of 81 coefficients. With two possible detections per segment, 3 values for position regression within each segment and one value for the probability, this results in a total of $2\\cdot(81+3+1)=170$ parameters per segment and, thus, per output voxel.<br>\n",
    "In the following a pre-trained network is used to predict those masks based on input images. Afterwards, predicted detections and shape are merged and post-processed to form the final predicted voxelized mask."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/temp/eschweiler/miniconda3/envs/tensorflowgpu/lib/python3.6/site-packages/tensorflow/python/framework/op_def_library.py:263: colocate_with (from tensorflow.python.framework.ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Colocations handled automatically by placer.\n"
     ]
    }
   ],
   "source": [
    "# Load the model architecture and the pre-trained weights\n",
    "with open('models/model.json','r') as mh:\n",
    "    model = model_from_json(mh.read())\n",
    "model.load_weights('models/model.hdf5')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Normalize the patch\n",
    "img_norm = img/255\n",
    "\n",
    "# Predict the encoding maskarchitecturearchitecture\n",
    "predicted_encoding = model.predict(img_norm[np.newaxis,..., np.newaxis])\n",
    "\n",
    "# Filter and decode the mask into the radii sampling\n",
    "# Since, for demonstration purposes, only one patch is processed, positional weighting of each detection\n",
    "# is disabled and the threshold values are slightly adapted to not exclude possibly truncated shape \n",
    "# predictions at the patch borders. \n",
    "# Usually a suitable overlap of neigbouring patches is used to only consider certain segmentations.\n",
    "pred_indices, pred_probs, pred_shapes = harmonicmask2sampling(predicted_encoding[0,...], \\\n",
    "                                                              h2s_converter=h2s_converter, \\\n",
    "                                                              cell_size=(8,8,8), \\\n",
    "                                                              dets_per_region=2, \\\n",
    "                                                              thresh=0.5, \\\n",
    "                                                              convert2radii=True, \\\n",
    "                                                              positional_weighting=False)\n",
    "\n",
    "# Perform clustering (used for each patch individually to reduce further computation)\n",
    "pred_indices, pred_probs, pred_shapes = agglomerative_clustering(pred_indices, pred_probs, shape_descriptors=pred_shapes, max_dist=10)\n",
    "\n",
    "# Perform NMS (Usually used after each patch has been processed and after they are merged to the full-size image)\n",
    "# Note that this is kind of overkill for a single patch, but reaches its full potential for merged full-size images.\n",
    "pred_indices, pred_probs, pred_shapes = harmonic_non_max_suppression(pred_indices, pred_probs, pred_shapes, overlap_thresh=0.5)\n",
    "\n",
    "# Reconstruct the voxelized segmentation mask\n",
    "pred_mask = descriptors2image([pred_indices, pred_probs, pred_shapes], theta_phi_sampling=theta_phi_sampling, shape=(64,64,64), thresh=0.5)        \n",
    "                        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 28520 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Plot the image and mask\n",
    "plt.subplot(1,2,1)\n",
    "plt.imshow(img[...,32])\n",
    "plt.subplot(1,2,2)\n",
    "plt.imshow(pred_mask[...,32])\n",
    "plt.show()"
   ]
  }
 ],
 "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.6.7"
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
