{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1c5180cb",
   "metadata": {},
   "source": [
    "In this practical session, our goal is to use the ToMATo algorithm to compute segmentations of 3D shapes, i.e., to assign labels to 3D shape vertices in an unsupervised way, that is, without training on known labels. This task was initially explored in [this article](https://www.lix.polytechnique.fr/~maks/papers/pers_seg.pdf). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e265352a",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib notebook\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a488f1e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import numpy as np\n",
    "import scipy as sp\n",
    "import matplotlib.pyplot as plt\n",
    "import sklearn.preprocessing as skp\n",
    "import gudhi as gd\n",
    "import gudhi.clustering.tomato as gdt\n",
    "import gudhi.representations   as gdr\n",
    "import meshplot as mp\n",
    "import statistics\n",
    "import robust_laplacian as rlap"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9475e32b",
   "metadata": {},
   "source": [
    "First things first, we have to download the data set. It can be obtained [here](https://people.cs.umass.edu/~kalo/papers/LabelMeshes/labeledDb.7z). Extract it, and save its path in the `dataset_path` variable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9cd1947",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_path = './3dshapes/'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6fd00c44",
   "metadata": {},
   "source": [
    "As you can see, the data set in split in several categories (`Airplane`, `Human`, `Teddy`, etc), each category having its own folder. Inside each folder, some 3D shapes (i.e., 3D triangulations) are provided in [`.off`](https://en.wikipedia.org/wiki/OFF_(file_format)) format, and face (i.e., triangle) labels are provided in text files (extension `.txt`). "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b10ffa39",
   "metadata": {},
   "source": [
    "Every data science project begins by some preprocessing ;-) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7fad42e",
   "metadata": {},
   "source": [
    "Write a function `off2numpy` that reads information from an `.off` file and store it in two `NumPy` arrays, called `vertices` (type float and shape number_of_vertices x 3---the 3D coordinates of the vertices) and `faces` (type integer and shape number_of_faces x 3---the IDs of the vertices that create faces). Write also a function `get_labels` that stores the face labels of a given 3D shape in a `NumPy` array (type string or integer and shape [number_of_faces].  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bcd76b9e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def off2numpy(shape_name):\n",
    "    with open(shape_name, 'r') as S:\n",
    "        S.readline()\n",
    "        num_vertices, num_faces, _ = [int(n) for n in S.readline().split(' ')]\n",
    "        info = S.readlines()\n",
    "    vertices = np.array([[float(coord) for coord in l.split(' ')] for l in info[0:num_vertices]])\n",
    "    faces    = np.array([[int(coord) for coord in l.split(' ')[1:]] for l in info[num_vertices:]])\n",
    "    return vertices, faces"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ddb31d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_labels(label_name, num_faces):\n",
    "    L = np.empty([num_faces], dtype='|S100')\n",
    "    with open(label_name, 'r') as S:\n",
    "        info = S.readlines()\n",
    "    labels, face_indices = info[0::2], info[1::2]\n",
    "    for ilab, lab in enumerate(labels):\n",
    "        indices = [int(f)-1 for f in face_indices[ilab].split(' ')[:-1]]\n",
    "        L[  np.array(indices)  ] = lab[:-1]\n",
    "    return L"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "09159d0d",
   "metadata": {},
   "source": [
    "You can now apply your code and use `meshplot` to visualize a given 3D shape, say `61.off` in `Airplane`, and the labels on its faces."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a7b3b5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "vertices, faces = off2numpy(dataset_path + 'Airplane/61.off')\n",
    "label_faces = get_labels(dataset_path + 'Airplane/61_labels.txt', len(faces))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e01afb1c",
   "metadata": {},
   "outputs": [],
   "source": [
    "mp.plot(vertices, faces, c=skp.LabelEncoder().fit_transform(label_faces))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6ccd696",
   "metadata": {},
   "source": [
    "If `meshplot` does not work, we also provide a fix with `matplotlib` (it requires converting the face labels into point labels though)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a841e9a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def face2points(vals_faces, faces, num_vertices):\n",
    "    vals_points = np.empty([num_vertices], dtype=type(vals_faces))\n",
    "    for iface, face in enumerate(faces):\n",
    "        vals_points[face] = vals_faces[iface]\n",
    "    return vals_points"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a873709e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def set_axes_equal(ax: plt.Axes):\n",
    "    limits = np.array([\n",
    "        ax.get_xlim3d(),\n",
    "        ax.get_ylim3d(),\n",
    "        ax.get_zlim3d(),\n",
    "    ])\n",
    "    origin = np.mean(limits, axis=1)\n",
    "    radius = 0.5 * np.max(np.abs(limits[:, 1] - limits[:, 0]))\n",
    "    _set_axes_radius(ax, origin, radius)\n",
    "\n",
    "def _set_axes_radius(ax, origin, radius):\n",
    "    x, y, z = origin\n",
    "    ax.set_xlim3d([x - radius, x + radius])\n",
    "    ax.set_ylim3d([y - radius, y + radius])\n",
    "    ax.set_zlim3d([z - radius, z + radius])\n",
    "    \n",
    "fig = plt.figure()\n",
    "ax = fig.add_subplot(projection='3d')\n",
    "ax.set_aspect('equal')\n",
    "ax.scatter(vertices[:,0], vertices[:,1], vertices[:,2], s=1, \n",
    "           c=skp.LabelEncoder().fit_transform(\n",
    "    face2points(label_faces, faces, len(vertices))))\n",
    "set_axes_equal(ax)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54460973",
   "metadata": {},
   "source": [
    "Overall, the main idea is to run ToMATo on the neighborhood graph given by the triangulation, with the so-called Heat Kernel Signature (HKS) as the filter. This is motivated by the fact that the HKS function typically takes higher values on the parts of the 3D shape that are very curved (such as, e.g., the tips of fingers in human hand shapes).  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c226c83",
   "metadata": {},
   "source": [
    "The HKS was defined in [this article](https://onlinelibrary.wiley.com/doi/epdf/10.1111/j.1467-8659.2009.01515.x). It is related to the heat equation on a given 3D shape $S$:\n",
    "\n",
    "$$\\Delta_S f = -\\frac{\\partial f}{\\partial t}.$$\n",
    "\n",
    "More formally, the HKS function  with parameter $t >0$ on a vertex $v\\in S$, and denoted by ${\\rm HKS}_t(v)$, is computed as:\n",
    "\n",
    "$${\\rm HKS}_t(v) = \\sum_{i=0}^{+\\infty} {\\rm exp}(-\\lambda_i\\cdot t)\\cdot \\phi_i^2(v),$$\n",
    "\n",
    "where $\\{\\lambda_i, \\phi_i\\}_i$ are the eigenvalues and eigenvectors of $\\Delta_S$.\n",
    "Intuitively, ${\\rm HKS}_t(v)$ is the amount of heat remaining on $v$ at time $t$, after unit sources of heat have been placed on each vertex at time `t=0`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f8faae24",
   "metadata": {},
   "source": [
    "Let's first pick a 3D shape. For instance, use `Hand/181.off` (or any other one you would like to try)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78a0a788",
   "metadata": {},
   "outputs": [],
   "source": [
    "vertices, faces = off2numpy(dataset_path + 'Hand/181.off')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "752f3cb1",
   "metadata": {},
   "source": [
    "Now, use `robust_laplacian` to compute the first 200 eigenvalues and eigenvectors of its Laplacian (you can use the `eigsh` function of `SciPy` for diagonalizing the Laplacian)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58400252",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "f041b182",
   "metadata": {},
   "source": [
    "Write a function `HKS` that uses these eigenvalues and eigenvectors, as well as a time parameter, to compute the HKS value on a given vertex."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00cf5fc5",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "260833bf",
   "metadata": {},
   "source": [
    "Visualize the function values with `meshplot` for different time parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cab40e8c",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "74dafe52",
   "metadata": {},
   "source": [
    "Recall that ToMATo requires, in addition to the filter, a neighborhood graph built on top of the data. Fortunately, we can use the triangulations of our 3D shapes as input graphs! Write a function `get_neighborhood_graph_from_faces` that computes a neighborhood graph (in the format required by ToMATo) from the faces of a triangulation. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bede360d",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "fd5e2ab6",
   "metadata": {},
   "source": [
    "Finally, apply ToMATo (with no prior on the number of clusters or merging threshold) on the neighborhood graph and the HKS function associated to a given time parameter."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "948698bb",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f6c6adc",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "4e7b7a1e",
   "metadata": {},
   "source": [
    "Visualize the persistence diagram produced by ToMATo."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4708b364",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "45129195",
   "metadata": {},
   "source": [
    "How many points do you see standing out from the diagonal? Use this number to re-cluster."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d2f9e78",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "6a3ae1b4",
   "metadata": {},
   "source": [
    "Visualize the 3D shape with the ToMATo labels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dfcfc95f",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "e2c6cbf0",
   "metadata": {},
   "source": [
    "Does our segmentation make sense? Can you interpret the boundaries between labels?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d7b9746",
   "metadata": {},
   "source": [
    "Since the boundaries are driven by the elder rule, they can seem a bit shaggy. In order to fix this, we can use bootstrap-like smoothing. The idea is to first save the current ToMATo clustering obtained with filter $f$ (let's call it the initial clustering), and then perturb $f$ a little bit into another function $\\tilde f$, and finally recompute clustering with ToMATo using $\\tilde f$. Since clusters are now created with the maxima of $\\tilde f$ (which will be different in general from those of $f$), we can use the initial clustering to relate the clusters of $\\tilde f$ to those of $f$, by simply looking at which (initial) clusters do the maxima of $\\tilde f$ belong to. If we repeat this procedure $N$ times, we will end up with a distribution (of size $N$) of candidate clusters for each vertex $v$. It suffices to pick the most frequent one for each vertex to get a smooth segmentation for the 3D shape. See also Section 6 in [the article](https://www.lix.polytechnique.fr/~maks/papers/pers_seg.pdf)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52af40db",
   "metadata": {},
   "source": [
    "In order to implement this, write first a function `get_indices_of_maxima` which computes the indices of the maxima  associated to a set of ToMATo clusters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb542b81",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "81323c64",
   "metadata": {},
   "source": [
    "Compute and plot these maxima on the 3D shape to make sure your code is working."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c6cf376",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "294d2ba0",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "1040afb7",
   "metadata": {},
   "source": [
    "Now, use this function to write another function `bootstrap_tomato` that perform a bootstrap smoothing of a set to ToMATo labels. This function will also take as arguments a number $N$ of bootstrap iterations, and a parameter $\\epsilon$ controlling the amplitude of the uniform noise used to perturb the filter."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "792f6bf5",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "60587ff8",
   "metadata": {},
   "source": [
    "Apply the bootstrap smoothing and visualize the segmentation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a451e08",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98819417",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "6df187a5",
   "metadata": {},
   "source": [
    "Is the segmentation any better? How does the result depend on the noise amplitude?"
   ]
  }
 ],
 "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.11.4"
  },
  "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": 5
}