{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Demonstration of the Synthesis Pipeline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generating Synthetic Membrane Masks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from skimage import io\n",
    "from scipy.ndimage import distance_transform_edt\n",
    "from util.synthetic_cell_membrane_masks import SyntheticCellMembraneMasks\n",
    "\n",
    "plt.rc('font', size=18)         # controls default text sizes\n",
    "plt.rc('axes', titlesize=12)    # fontsize of the axes title\n",
    "plt.rc('axes', labelsize=12)    # fontsize of the x and y labels\n",
    "plt.rc('xtick', labelsize=12)   # fontsize of the tick labels\n",
    "plt.rc('ytick', labelsize=12)   # fontsize of the tick labels\n",
    "plt.rc('figure', titlesize=32)  # fontsize of the figure title"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before we can generate a synthetic mask, some parameters have to be specified. For demonstration purposes, the simulation is limited to a small specimen and, thus, a small number of cells.\n",
    "Cell density is estimated from the training set and was determined to be roughly one cell per $20\\times20\\times20$ pixel region."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Set parameters\n",
    "\n",
    "img_size = 128 # So far only quadratic images can be created\n",
    "cell_density = 1/(20*20*20)\n",
    "roundness_gamma = 5\n",
    "sphere_radius = img_size//2\n",
    "cell_count = np.int(4/3+np.pi*(sphere_radius**3)*cell_density)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Set up the generator\n",
    "\n",
    "# Any time this generator is called, a randomly shaped specimen is generated\n",
    "membrane_generator = SyntheticCellMembraneMasks(size=img_size, radius=sphere_radius, roundness=roundness_gamma, cell_count=cell_count, sphere_count=1, use_cuda=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 13520 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Plot 2D crops of the results\n",
    "\n",
    "plt.subplot(121)\n",
    "plt.title('Instance mask')\n",
    "plt.imshow(membrane_generator.label_image[img_size//2,:,:])\n",
    "plt.subplot(122)\n",
    "plt.title('Boundary of the instance mask')\n",
    "plt.imshow(membrane_generator.labelBoundaries(membrane_generator.label_image)[img_size//2,:,:])\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As exlpained in the paper, before generating the final instance mask, boundary morphology is enhanced by applying morphological opening. The change can be observed in the following plot:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 12656 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.subplot(121)\n",
    "plt.title('Initial boundary mask')\n",
    "plt.imshow(membrane_generator.boundary_image[img_size//2,:,:])\n",
    "plt.subplot(122)\n",
    "plt.title('Enhanced boundary mask')\n",
    "plt.imshow(membrane_generator.labelBoundaries(membrane_generator.label_image)[img_size//2,:,:])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data Preparation\n",
    "\n",
    "As input for the generative adversarial network, two distance maps have to be created. Those maps encode the position at each pixel by a tuple of distances from the outer boundary of the specimen to the centroid and to the background region, respectively. The resulting 3-channel image serves as input for the generator network."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 28264 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Get membrane mask and background label\n",
    "instance_mask = membrane_generator.label_image\n",
    "membrane_mask = membrane_generator.labelBoundaries(instance_mask)\n",
    "background_label = np.unique(membrane_mask)[0]\n",
    "\n",
    "# Calculate distance maps\n",
    "dist_map1 = distance_transform_edt(instance_mask==background_label)\n",
    "dist_map1 = np.max(dist_map1)-dist_map1 # Invert the distances\n",
    "dist_map1[instance_mask!=background_label] = 0 # Remove foreground distances\n",
    "dist_map2 = distance_transform_edt(instance_mask!=background_label)\n",
    "dist_map2 = np.max(dist_map2)-dist_map2 # Invert the distances\n",
    "dist_map2[instance_mask==background_label] = 0 # Remove background distances\n",
    "\n",
    "# Concatenate the image\n",
    "membrane_patch = np.concatenate((membrane_mask[np.newaxis,np.newaxis,...],\\\n",
    "                              dist_map1[np.newaxis,np.newaxis,...],\\\n",
    "                              dist_map2[np.newaxis,np.newaxis,...]), axis=1)\n",
    "\n",
    "\n",
    "# Plot 2D slices\n",
    "plt.subplot(131)\n",
    "plt.title('Boundary mask')\n",
    "plt.imshow(membrane_generator.labelBoundaries(membrane_mask)[img_size//2,:,:])\n",
    "plt.subplot(132)\n",
    "plt.title('Distance map 1')\n",
    "plt.imshow(dist_map1[img_size//2,:,:])\n",
    "plt.subplot(133)\n",
    "plt.title('Distance map 2')\n",
    "plt.imshow(dist_map2[img_size//2,:,:])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generating Synthetic Microscopy Images\n",
    "Those 3-channel images are translated into the image domain by a generative adversarial network."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "initialize network with normal\n",
      "loading the model from /home/staff/eschweiler/Publications/CVPR2020_HarmonicSegmentationSynthesis/Supplementary/Synthesis/models/latest_net_G.pth\n"
     ]
    }
   ],
   "source": [
    "# Include the model class\n",
    "from models.gan_3D_model import GAN3D\n",
    "from util.model_setup import load_config_json, numpy_to_torch_tensor\n",
    "\n",
    "# Set up a config file\n",
    "general_config, image_config, mask_config = load_config_json('models/gan_patchwise.json', globals())\n",
    "\n",
    "# Set up the model and load pre-trained weights\n",
    "model_synthesis = GAN3D(general_config)\n",
    "model_synthesis.load_networks('latest')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Since the generator network was trained to work in much bigger specimen,\n",
    "# and the generation of large images would take too long for demonstration,\n",
    "# an examplary image from the public data set is used\n",
    "\n",
    "test_img = io.imread('demo_image.tif')\n",
    "test_img = np.transpose(test_img, (2,1,0))\n",
    "\n",
    "test_membranes = io.imread('demo_membrane.tif')\n",
    "test_membranes = np.transpose(test_membranes, (2,1,0))\n",
    "\n",
    "test_dist1 = io.imread('demo_dist1.tif')\n",
    "test_dist1 = np.transpose(test_dist1, (2,1,0))\n",
    "\n",
    "test_dist2 = io.imread('demo_dist2.tif')\n",
    "test_dist2 = np.transpose(test_dist2, (2,1,0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 69392 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Construct the input image\n",
    "inputA = test_img[np.newaxis,np.newaxis,...]/255\n",
    "\n",
    "inputB = np.concatenate((test_membranes[np.newaxis,np.newaxis,...],\\\n",
    "                         test_dist1[np.newaxis,np.newaxis,...],\\\n",
    "                         test_dist2[np.newaxis,np.newaxis,...]), axis=1)\n",
    "\n",
    "# Plot the patch\n",
    "plt.subplot(121)\n",
    "plt.title('Real Patch')\n",
    "plt.imshow(inputA[0,0,:,:,50])\n",
    "plt.subplot(122)\n",
    "plt.title('Real Mask')\n",
    "plt.imshow(inputB[0,0,:,:,50])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Predict the corresponding synthetic image \n",
    "model_synthesis.set_input([numpy_to_torch_tensor([inputA])[0], numpy_to_torch_tensor([inputB])[0]])\n",
    "model_synthesis.test()\n",
    "images = model_synthesis.get_current_visuals()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 115648 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Plot the generated image\n",
    "# For the full-size image, a faded overlap technique and border cropping is used to \n",
    "# reduce border actifacts and obtain continuous intensities.\n",
    "plt.subplot(121)\n",
    "plt.title('Real Patch')\n",
    "plt.imshow(images['real_B'].detach().cpu().numpy()[0,0,:,:,50])\n",
    "plt.subplot(122)\n",
    "plt.title('Fake Patch')\n",
    "plt.imshow(images['fake_B'].detach().cpu().numpy()[0,0,:,:,50])\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.7.0"
  },
  "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
}
