{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Getting Started with Data Augmentation"
   ]
  },
  {
   "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",
    "- Note example output was generated on Linux/Mac environment and may look different on Windows.\n",
    "\n",
    "## In this notebook, you will learn:\n",
    "\n",
    "1. How to generate realistic synthetic data from an existing dataset using different parametric distributions.\n",
    "2. How to visualize the statistical distribution of the generated data compared to the original data.\n",
    "\n",
    "### Data Augmentation Overview\n",
    "\n",
    "ShapeWorks includes a Python package, DataAugmentationUtils, that supports model-based data augmentation. This package is useful to increase the training sample size to train deep networks such as DeepSSM (see [SSMs Directly from Images](http://sciinstitute.github.io/ShapeWorks/deep-learning/deep-ssm.html)).\n",
    "\n",
    "A preliminary requirement for data augmentation is a set of images and shape models from real data on which to base augmentation. Once that is acquired, the process includes:\n",
    "1. Embedding the real data into a low-dimensional space using principle component analysis (PCA).\n",
    "2. Fitting a parametric distribution to the subspace for sampling.\n",
    "3. Sampling from the distribution to create new instances.\n",
    "4. Projecting the samples back into the high-dimensional space of the original data \n",
    "5. Completing the sample generation by creating a corresponding synthetic image. \n",
    "\n",
    "This notebook shows how the distribution of the original data can be visually compared to the distribution of the synthetic data to motivate the choice of parametric distribution in step 2. \n",
    "\n",
    "For a full explanation of the data augmentation process and package please see: [Data Augmentation for Deep Learning](http://sciinstitute.github.io/ShapeWorks/deep-learning/data-augmentation.html)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Import shapeworks and relevant libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os \n",
    "import sys \n",
    "import shapeworks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Import Data Augmentation Package"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import DataAugmentationUtils"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Defining the original dataset\n",
    "\n",
    "### Defining dataset location\n",
    "\n",
    "You can download exemplar datasets from [ShapeWorks data portal](https://girder.shapeworks-cloud.org) after you login. For new users, you can [register](https://girder.shapeworks-cloud.org/#?dialog=register) an account for free. Please do not use an important password.\n",
    "\n",
    "After you login, click `Collections` on the left panel and then `use-case-data-v2`. Select the dataset you would like to download by clicking on the checkbox on the left of the dataset name. See the video below.\n",
    "After you download the dataset zip file, make sure you unzip/extract the contents in the appropriate location.\n",
    "\n",
    "**This notebook assumes that you have downloaded `femur-v0` and you have placed the unzipped folder `femur-v0` in `Examples/Python/Data`.** Feel free to use your own dataset. \n",
    "\n",
    "\n",
    "<p><video src=\"https://sci.utah.edu/~shapeworks/doc-resources/mp4s/portal_data_download.mp4\" autoplay muted loop controls style=\"width:100%\"></p>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# dataset name is the folder name for your dataset\n",
    "datasetName  = 'femur-v0'\n",
    "\n",
    "# path to the dataset where we can find shape data \n",
    "# here we assume shape data are given as binary segmentations\n",
    "data_dir      = '../../Data/' + datasetName + '/'\n",
    "    \n",
    "print('Dataset Name:     ' + datasetName)\n",
    "print('Directory:  ' + data_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Get file lists\n",
    "Now we need the .particle files and corresponding raw images for the original dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get image path list\n",
    "img_dir = data_dir + \"groomed/images/\"\n",
    "img_list = []\n",
    "for file in os.listdir(img_dir):\n",
    "    img_list.append(img_dir + file)\n",
    "img_list = sorted(img_list)\n",
    "\n",
    "# Get particles path list\n",
    "model_dir =  data_dir + \"shape_models/femur/1024/\" \n",
    "world_particle_list = []\n",
    "for file in os.listdir(model_dir):\n",
    "    if \"world\" in file:\n",
    "        world_particle_list.append(model_dir + file)\n",
    "world_particle_list = sorted(world_particle_list)\n",
    "\n",
    "print(\"Total shapes in original dataset: \"+ str(len(img_list)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run data augmentation using a Gaussian Distribution\n",
    "\n",
    "Below is the command for running the complete data augmentation process:\n",
    "\n",
    "```python\n",
    "DataAugmentationUtils.runDataAugmentation(out_dir, img_list, \n",
    "                                          world_point_list, num_samples, \n",
    "                                          num_dim, percent_variability, \n",
    "                                          sampler_type, mixture_num)\n",
    "```\n",
    "**Input arguments:**\n",
    "\n",
    "* `out_dir`: Path to the directory where augmented data will be stored\n",
    "* `img_list`: List of paths to images of the original dataset.\n",
    "* `world_point_list`: List of paths to world `.particles` files of the original dataset. Note, this list should be ordered in correspondence with the `img_list`.\n",
    "* `num_dim`: The number of dimensions to reduce to in PCA embedding. If zero or not specified, the percent_variability option is used to select the numnber of dimensions.\n",
    "* `percent_variability`: The proportion of variability in the data to be preserved in embedding. Used if `num_dim` is zero or not specified. Default value is 0.95 which preserves 95% of the varibaility in the data.\n",
    "* `sampler_type`: The type of parametric distribution to fit and sample from. Options: `gaussian`, `mixture`, or `kde`. Default: `kde`.\n",
    "* `mixture_num`: Only necessary if `sampler_type` is `mixture`. The number of clusters (i.e., mixture components) to be used in fitting a mixture model. If zero or not specified, the optimal number of clusters will be automatically determined using the [elbow method](https://en.wikipedia.org/wiki/Elbow_method_(clustering)).\n",
    "\n",
    "\n",
    "In this notebook we will keep most arguments the same and explore the effect of changing the `sampler_type`.\n",
    "First, we will try a Gaussian distribution. For further explanation about each distribution, see [Data Augmentation for Deep Learning](http://sciinstitute.github.io/ShapeWorks/deep-learning/data-augmentation.html). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Augmentation variables to keep constant\n",
    "num_samples = 50\n",
    "num_dim = 0\n",
    "percent_variability = 0.95"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_directory = '../Output/GaussianAugmentation/'\n",
    "sampler_type = \"gaussian\"\n",
    "embedded_dim = DataAugmentationUtils.runDataAugmentation(output_directory, img_list, world_particle_list, num_samples, num_dim, percent_variability, sampler_type)\n",
    "aug_data_csv = output_directory + \"/TotalData.csv\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualize distribution of real and augmented data\n",
    "\n",
    "Below is the command for visualizing the original and augmented data:\n",
    "```\n",
    "DataAugmentationUtils.visualizeAugmentation(data_csv, viz_type)\n",
    "```\n",
    "**Input arguments:**\n",
    "\n",
    "* `data_csv`: The path to the CSV file created by running the data augmentation process.\n",
    "* `viz_type`: The type of visulazation to display. Options `splom` or `violin` (default: `splom`). If set to `splom`, a scatterplot matrix of pairwise PCA comparisions will open in the default browser. If set to `violin` a violin plot or rotated kernel density plot will be displayed. \n",
    "\n",
    "We will use a violin plot to visualize the difference in the real and augmented distributions. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DataAugmentationUtils.visualizeAugmentation(aug_data_csv, 'violin')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Example output:\n",
    "<p><img src=\"https://sci.utah.edu/~shapeworks/doc-resources/pngs/data_aug_gaussian.png\"></p>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run data augmentation using a Mixture of Gaussian Distribution\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_directory = '../Output/MixtureAugmentation/'\n",
    "sampler_type = \"mixture\"\n",
    "embedded_dim = DataAugmentationUtils.runDataAugmentation(output_directory, img_list, world_particle_list, num_samples, num_dim, percent_variability, sampler_type)\n",
    "aug_data_csv = output_directory + \"/TotalData.csv\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualize distribution of real and augmented data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DataAugmentationUtils.visualizeAugmentation(aug_data_csv, 'violin')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Example output:\n",
    "<p><img src=\"https://sci.utah.edu/~shapeworks/doc-resources/pngs/data_aug_mixture.png\"></p>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run data augmentation using Kernel Density Estimation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_directory = '../Output/KDEAugmentation/'\n",
    "sampler_type = \"kde\"\n",
    "embedded_dim = DataAugmentationUtils.runDataAugmentation(output_directory, img_list, world_particle_list, num_samples, num_dim, percent_variability, sampler_type)\n",
    "aug_data_csv = output_directory + \"/TotalData.csv\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualize distribution of real and augmented data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DataAugmentationUtils.visualizeAugmentation(aug_data_csv, 'violin')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Example output:\n",
    "<p><img src=\"https://sci.utah.edu/~shapeworks/doc-resources/pngs/data_aug_kde.png\"></p>"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.13"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "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": 1
}
