{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Getting Started with Shape Cohort Generator"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Before you start!\n",
    "\n",
    "- This notebook assumes that shapeworks conda environment has been activated using `conda activate shapeworks` on the terminal.\n",
    "- See [Getting Started with Notebooks](getting-started-with-notebooks.ipynb) for information on setting up your environment and notebook shortcuts. \n",
    "- See [Getting Started with Meshes](getting-started-with-meshes.ipynb) to learn how to load and visualize meshes.\n",
    "- See [Getting Started with Segmentations](getting-started-with-segmentations.ipynb) to learn how to load and visualize binary segmentations.\n",
    "- Note example output was generated on Linux/Mac environment and may look different on Windows.\n",
    "\n",
    "\n",
    "## In this notebook, you will learn:\n",
    "\n",
    "\n",
    "How to use the `ShapeCohortGenerator` package to generate meshes and segmentations (binary images) for synthetic shape cohorts, i.e., parameterized families of shapes. \n",
    "\n",
    "\n",
    "\n",
    "## About `ShapeCohortGenerator`\n",
    "\n",
    "`ShapeCohortGenerator` is a python package that generates synthetic shape cohorts with groundtruth surface correspondences by varying different parameters describing such shape families.\n",
    "\n",
    "\n",
    "### What is a shape cohort ?\n",
    "\n",
    "A shape cohort is a collection of geometric shapes that attain clear differences in shape; however, they share common characteristics that stem from the underlying mechanisms involved in their formation. For real-world shapes, e.g., anatomical structures, such common characteristics (or factor of variations) are not known in advance, hence ShapeWorks discovers such factors of variations directly from surface meshes or binary segmentations of such shapes. `ShapeCohortGenerator` uses the true factors of variations known for synthetic shapes that are analytically parameterized.\n",
    "\n",
    "\n",
    "### Why `ShapeCohortGenerator`?\n",
    "\n",
    "We require a shape population dataset to run the shape modeling workflow. Each population dataset requires unique grooming steps. Developing and testing complicated grooming pipelines for large-scale datasets can consume a lot of computational resources and time. Hence, having a few toy datasets, which are lightweight and robust in variability can make this development and debugging process easier and simpler.These cohorts can also be used to test the optimization workflow. \n",
    "\n",
    "\n",
    "### What families of shape can be generated by `ShapeCohortGenerator`?\n",
    "\n",
    "`ShapeCohortGenerator` currently supports two families of synthetic shapes, namely [ellipsoids](https://en.wikipedia.org/wiki/Ellipsoid) and  [supershapes](https://link.springer.com/article/10.1007/s00158-018-2034-z?shared-article-renderer).  \n",
    "\n",
    "\n",
    "#### Ellipsoids\n",
    "An [ellipsoid](https://en.wikipedia.org/wiki/Ellipsoid)is symmetrical about three mutually perpendicular axes that intersect at the center. If a, b, and c are the principal semiaxes, the general equation of such an ellipsoid is \n",
    "   $$\\frac{x^2}{a^2} + \\frac{y^2}{b^2} + \\frac{z^2}{c^2} = 1$$\n",
    "   \n",
    "#### Supershapes\n",
    "[Supershapes](https://link.springer.com/article/10.1007/s00158-018-2034-z?shared-article-renderer) are an extension of [superellipses](https://pubmed.ncbi.nlm.nih.gov/21659124/) that can exhibit variable symmetry as well as asymmetry. Supershapes can be described through a single equation, the so-called superformula, that parametrizes a wide variety of shapes, including geometric primitives. The superformula is given by :\n",
    "\n",
    "$$ r(\\theta) = \\left[ \\left| \\frac{1}{a} \\cos \\left( \\frac{m\\theta}{4} \\right) \\right|^{n_2} +  \\left|  \\frac{1}{b} \\sin \\left(\\frac{m\\theta}{4} \\right) \\right|^{n_3}  \\right]^{-\\frac{1}{n_1}} $$\n",
    "    \n",
    "Unlike superellipses, supershapes need not to be symmetric; the parameter $m$ controls the rotational symmetry. The values of  $a$ and $b$ control the size, and the exponents $n_1,n_2$ and $n_3$ control the curvature of the sides. The superformula can produce a wide range of shapes,including many shapes found in nature. \n",
    "\n",
    "The `ShapeCohortGenerator` package allows the user to specify the rotational symmetry $m$ and the size. The values of $n_1,n_2$ and $n_3$ are randomly selected to creates shapes with different curvatures. Examples of these supershapes with different $m$ values can be seen below.\n",
    "\n",
    "<p><img src=\"https://sci.utah.edu/~shapeworks/doc-resources/gifs/supershapes_examples.png\"></p>\n",
    "\n",
    "#### 2D Supershapes\n",
    "2D supershapes are the same as Supershapes but two-dimensional instead of three. When generation is called, contours are created rather than meshes When generate_segmentations and generate_images are called, 2D .png images are created rather than 3D .nrrd images. \n",
    "\n",
    "#### Joint Ellipsoids \n",
    "Joint ellipsoids are two ellipsoids positioned in a parameterized way in relation to each other. This cohort allows for exploring modes of variation such as relative rotation.\n",
    "\n",
    "#### Torus\n",
    "Tori are donut-shaped surfaces generated by a circle rotated about an axis in its plane that does not intersect the circle. For the tori mesh generation, the arguments are the same as for ellipsoids, except instead of `randomize_x_radius` and `randomize_y_radius` we have `randomize_ring_radius` (outside) and `randomize_cross_section_radius`(inside).\n",
    "<p><img src=\"https://sci.utah.edu/~shapeworks/doc-resources/pngs/tori.png\"></p>\n",
    "\n",
    "### What you can do with `ShapeCohortGenerator`?\n",
    "\n",
    "The `ShapeCohortGenerator` package can be used to generate collections of ellipsoids or supershapes, where the user can control the number of shapes in the cohort and the variability of the members of the cohort.\n",
    "\n",
    "Each cohort will have mesh data (vtk format, supershapes will be saved in stl format) and segmentation image data (nrrd format). These cohorts generated by the package can be directly run with `ShapeWorks`. Generating these cohorts in the `Output` folder would be a good way to start.\n",
    "\n",
    "\n",
    "We will also define modular/generic helper functions as we walk through these items to reuse functionalities without duplicating code."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Importing `shapeworks` library"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import shapeworks as sw"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Importing `ShapeCohortGen` library\n",
    "To use this package, first a `generator` is defined, then `generate()` is called that generates shapes in vtk mesh format. Then segmentations (binary image) and images (synthetic intensities that mimic imaging data for real shapes) can be created from those meshes.\n",
    "\n",
    "Each generator has three functions:\n",
    "\n",
    "- `generate()` for mesh generation (function specific to generator type)\n",
    "- `generate_segmentations()` for segmentation generation based on meshes (general function shared by all generator types)\n",
    "- `generate_images()` for image generation based on segmentations (general function shared by all generator types)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ShapeCohortGen"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Defining parameters for `pyvista` plotter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define parameters that controls the plotter\n",
    "\n",
    "# common for volumes and meshes visualization\n",
    "notebook       = False # True will set the renderings inline\n",
    "show_borders   = True  # show borders for each rendering window\n",
    "show_axes      = True  # show a vtk axes widget for each rendering window\n",
    "show_bounds    = True  # show volume bounding box\n",
    "show_all_edges = True  # add an unlabeled and unticked box at the boundaries of plot. \n",
    "font_size      = 10    # text font size for windows\n",
    "link_views     = True  # link all rendering windows so that they share same camera and axes boundaries\n",
    "\n",
    "# for volumes\n",
    "shade_volumes  = True  # use shading when performing volume rendering\n",
    "color_map       = 'coolwarm' # color map for volume rendering, e.g., 'bone', 'coolwarm', 'cool', 'viridis', 'magma'\n",
    "\n",
    "# for meshes\n",
    "meshes_color    = 'tan' # color to be used for meshes (can be a list with the same size as meshList if different colors are needed)\n",
    "mesh_style      = 'surface' # visualization style of the mesh. style='surface', style='wireframe', style='points'. \n",
    "show_mesh_edges = False  # show mesh edges "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generating an Ellipsoid Cohort\n",
    "\n",
    "### Step 1: Initalize Ellipsoid Generator\n",
    "\n",
    "Here, we will initialize an ellipsoid cohort generator. The output directory needs to be specified, otherwise an output directory will automatically generated.\n",
    "\n",
    "Arguments:\n",
    "- `out_dir`: path where the dataset should be saved<br>\n",
    "    Datatype : `string`<br> \n",
    "    Default value : `current_directory/generated_ellipsoid_cohort/` <br>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "out_dir             = \"../Output/Generated_Ellipsoids/\"\n",
    "ellipsoid_generator = ShapeCohortGen.EllipsoidCohortGenerator(out_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 2: Generate Meshes\n",
    "\n",
    "For the ellipsoid mesh generation, you can specify the following arguments:\n",
    "- `num_samples`: number of samples in the cohort(dataset)<br>\n",
    "    Datatype : `int` <br>\n",
    "    Default value : 3 <br>\n",
    "    <br>\n",
    "- `randomize_center`: randomizes the centers for ellipsoid mesh generation if set to `True`<br>\n",
    "    Datatype : `bool` <br> \n",
    "    Defaut value : `True` <br>\n",
    "    <br>   \n",
    "- `randomize_rotation`: randomizes the orientation of the ellipsoid if set to `True` <br>\n",
    "    Datatype : `bool` <br> \n",
    "    Defaut value : `True`\n",
    "    <br>\n",
    "    \n",
    "- `randomize_x_radius`: randomizes the radius of the ellipsoid along x-axis if set to `True` or else the value is fixed as 20 for all ellipsoids<br>\n",
    "    Datatype : `bool` <br> \n",
    "    Defaut value : `True`\n",
    "    <br>\n",
    "- `randomize_y_radius`: randomizes the radius of the ellipsoid along y-axis if set to `True` or else the value is fixed as 10 for all ellipsoids<br>\n",
    "    Datatype : `bool` <br> \n",
    "    Defaut value : `True`\n",
    "    <br>\n",
    "- `randomize_z_radius`: randomizes the radius of the ellipsoid along z-axis if set to `True` or else the value is fixed as 10 for all ellipsoids<br>\n",
    "    Datatype : `bool` <br> \n",
    "    Defaut value : `True`\n",
    "    <br>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_samples = 8\n",
    "meshFiles   = ellipsoid_generator.generate(num_samples)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "meshList = []\n",
    "for i in range(len(meshFiles)):\n",
    "    shapeMesh = sw.Mesh(meshFiles[i])\n",
    "    meshList.append(shapeMesh)\n",
    "\n",
    "# Plot the meshes\n",
    "sw.plot_meshes(meshList, use_same_window=True, notebook=notebook)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<p><img src=\"https://sci.utah.edu/~shapeworks/doc-resources/pngs/generated_ellipsoids.png\"></p>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 3: Generate Segmentations\n",
    "\n",
    "For segmentation generation, you can specify the following arguments:\n",
    "\n",
    "- `randomize_size`: randomize the size of the images to include more background if set to `True`<br>\n",
    "    Datatype : `bool` <br> \n",
    "    Defaut value : `True`\n",
    "    <br>\n",
    "    \n",
    "- `spacing`: set the spacing of the segmentation image <br>\n",
    "    Datatype: `list` <br>\n",
    "    Default value: `[1,1,1]` <br>\n",
    "    <br>\n",
    "- `allow_on_boundary`: If set to `True`,randomly selects 20% samples and ensure that the shapes are touching two random selected axes out of `[x,y,z]`<br>\n",
    "    Datatype : `bool` <br> \n",
    "    Defaut value : `True`\n",
    "    <br>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "segFiles = ellipsoid_generator.generate_segmentations()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's visualize the generated segmentations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "shapeSegList = []\n",
    "shapeNames   = []\n",
    "\n",
    "for segFile in segFiles:\n",
    "    shapeSegList.append(sw.Image(segFile))\n",
    "    shapeNames.append(segFile.split('/')[-1])\n",
    "print(shapeNames)\n",
    "\n",
    "sw.plot_volumes(shapeSegList,    \n",
    "             volumeNames    = shapeNames, \n",
    "             notebook       = notebook, \n",
    "             show_borders   = show_borders,  \n",
    "             shade_volumes  = shade_volumes,  \n",
    "             show_axes      = show_axes,  \n",
    "             show_bounds    = show_bounds,\n",
    "             show_all_edges = show_all_edges, \n",
    "             font_size      = font_size,\n",
    "             link_views     = True ) #link_views"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<p><img src=\"https://sci.utah.edu/~shapeworks/doc-resources/pngs/generated_ellipsoids_segs.png\"></p>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 4: Generate Images - Turning segmentations into non-binary images\n",
    "\n",
    "For the image generation, a Gaussian distribution is used to define foreground and background pixels values and a blur factor is used to blur the boundary with a Gaussian filter. You can specify the following arguments:\n",
    "\n",
    "- `blur_factor`: size of Gaussian filter to use for boundary blurring <br>\n",
    "    Datatype : `int` <br> \n",
    "    Defaut value : `1`\n",
    "    <br>\n",
    "- `foreground_mean`: mean of the foreground pixel value distribution <br>\n",
    "    Datatype: `int` <br>\n",
    "    Default value: `180` <br>\n",
    "    <br>\n",
    "- `foreground_var`: variance of the foreground pixel value distribution <br>\n",
    "    Datatype : `int` <br> \n",
    "    Defaut value : `30`\n",
    "    <br>\n",
    "- `background_mean`: mean of the background pixel value distribution <br>\n",
    "    Datatype: `int` <br>\n",
    "    Default value: `80` <br>\n",
    "    <br>\n",
    "- `background_var`: variance of the foreground pixel value distribution <br>\n",
    "    Datatype : `int` <br> \n",
    "    Defaut value : `30`\n",
    "    <br>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "imageFiles = ellipsoid_generator.generate_images()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's compare a segmentation to it's corresponding image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Segmentation:\")\n",
    "seg0 = sw.Image(segFiles[0])\n",
    "sw.plot_volumes(seg0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<p><img src=\"https://sci.utah.edu/~shapeworks/doc-resources/pngs/generated_ellipsoid_seg.png\"></p>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Image:\")\n",
    "img0 = sw.Image(imageFiles[0])\n",
    "sw.plot_volumes(img0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<p><img src=\"https://sci.utah.edu/~shapeworks/doc-resources/pngs/generated_ellipsoid_image.png\"></p>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generating Supershapes Cohort\n",
    "\n",
    "SuperShapes are parameterized shapes that have geometry based on a given number of lobes, $m$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 1: Initialize SuperShapes Generator\n",
    "\n",
    "Here, we will initialize  SuperShapes cohort generator. The output directory needs to be specified otherwise an output directory will automatically be generated.\n",
    "\n",
    "Argument:\n",
    "- `out_dir`: path where the dataset should be saved<br>\n",
    "    Datatype : `string`<br> \n",
    "    Default value : 'current_directory/generated_supershapes_cohort/' <br>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "out_dir      = \"../Output/Generated_Supershapes/\"\n",
    "ss_generator = ShapeCohortGen.SupershapesCohortGenerator(out_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 2: Generate Meshes\n",
    "\n",
    "For the supershapes mesh generation, you can specify the following arguments:\n",
    "- `num_samples` - number of samples in the cohort(dataset)<br>\n",
    "    Datatype : `int` <br>\n",
    "    Default value : 3 <br>\n",
    "    <br>\n",
    "- `randomize_center`: randomizes the centers for ellipsoid mesh generation if set to `True`<br>\n",
    "    Datatype : `bool` <br> \n",
    "    Defaut value : `True` <br>\n",
    "    <br>   \n",
    "- `randomize_rotation`: randomizes the orientation of the ellispoids if set to `True` <br>\n",
    "    Datatype : `bool` <br> \n",
    "    Defaut value : `True`\n",
    "    <br>\n",
    "- `m`: number of lobes supershapes should have <br>\n",
    "    Datatype : `int` <br>\n",
    "    Default value: `3` <br>\n",
    "    <br>\n",
    "- `size`: size of meshes (won't be more than 'size' away from center in any direction) <br>\n",
    "    Datatype: `int` <br>\n",
    "    Default value: `20` <br>\n",
    "    <br>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_samples = 8\n",
    "meshFiles   = ss_generator.generate(num_samples)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "meshList = []\n",
    "for i in range(len(meshFiles)):\n",
    "    shapeMesh = sw.Mesh(meshFiles[i])\n",
    "    meshList.append(shapeMesh)\n",
    "    \n",
    "# Plot the meshes\n",
    "sw.plot_meshes(meshList, use_same_window=True, notebook=notebook)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<p><img src=\"https://sci.utah.edu/~shapeworks/doc-resources/pngs/generated_supershapes.png\"></p>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 3: Generate Segmentations\n",
    "\n",
    "This is data type independent, the options are the same as they were for the ellipsoid."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "segFiles = ss_generator.generate_segmentations()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's visualize the generated segmentations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "shapeSegList = []\n",
    "shapeNames = []\n",
    "\n",
    "for segFile in segFiles:\n",
    "    shapeSegList.append(sw.Image(segFile))\n",
    "    shapeNames.append(segFile.split('/')[-1])\n",
    "print(shapeNames)\n",
    "\n",
    "sw.plot_volumes(shapeSegList,    \n",
    "             volumeNames    = shapeNames, \n",
    "             notebook       = notebook, \n",
    "             show_borders   = show_borders,  \n",
    "             shade_volumes  = shade_volumes,  \n",
    "             show_axes      = show_axes,  \n",
    "             show_bounds    = show_bounds,\n",
    "             show_all_edges = show_all_edges, \n",
    "             font_size      = font_size,   \n",
    "             link_views     = True ) #link_views)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<p><img src=\"https://sci.utah.edu/~shapeworks/doc-resources/pngs/generated_supershapes_segs.png\"></p>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 4: Generate Images\n",
    "\n",
    "This is also a standard function and has all the same options as listed before."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "imageFiles = ss_generator.generate_images()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's compare a segmentation to it's corresponding image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Segmentation:\")\n",
    "seg0 = sw.Image(segFiles[0])\n",
    "sw.plot_volumes(seg0, notebook=notebook)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<p><img src=\"https://sci.utah.edu/~shapeworks/doc-resources/pngs/generated_supershapes_seg.png\"></p>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Image:\")\n",
    "img0 = sw.Image(imageFiles[0])\n",
    "sw.plot_volumes(img0, notebook=notebook)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<p><img src=\"https://sci.utah.edu/~shapeworks/doc-resources/pngs/generated_supershapes_image.png\"></p>"
   ]
  }
 ],
 "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.8"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": false,
   "sideBar": true,
   "skip_h1_title": true,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
