main.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ce7b42c-311f-4318-9c21-e50fd8776542",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run this cell once to install everything, then run them all (including this one)\n",
    "import torch\n",
    "\n",
    "IN_COLAB = 'google.colab' in str(get_ipython()) # Check if we are in Google Colab or local\n",
    "if IN_COLAB :\n",
    "    try:\n",
    "        import mesh2sh\n",
    "    except ModuleNotFoundError:\n",
    "        !apt-get install libsuitesparse-dev && pip install scikit-sparse\n",
    "        !pip install git+git://github.com/le-Greg/mesh2sh.git\n",
    "        !pip install pyvista trimesh itkwidgets\n",
    "        print('The necessary packages have been installed. You have to restart the runtime to continue')\n",
    "        exit()\n",
    "    from google.colab import output\n",
    "    output.enable_custom_widget_manager()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c359bac0",
   "metadata": {},
   "source": [
    "## Load the mesh\n",
    "\n",
    "We use Trimesh to load, check and format the mesh. It is represented by its vertices, faces and colors. We place ourselves in the simplest configuration, where all the faces are triangular and the colors are defined for each vertex. The mesh must also be manifold and of genus 0.\n",
    "\n",
    "With V the number of vertices and F the number of faces,\n",
    "- The Verts array contains the XYZ coordinates for each vertex. Floats, shape \\[V, 3\\]\n",
    "- The Faces array contains the indices of the vertices connected in the original Verts array. Indexes are ordered counter-clockwise facing outward. Int64, shape \\[F, 3\\]\n",
    "- The Colors array contains the RGB colors of each vertex, in the same order as the Verts array. Colors range from 0 to 1. Float, shape \\[V, 3\\]\n",
    "\n",
    "All vectors are Pytorch vectors, allowing the GPU to be used to do the calculations. You can however use the CPU by changing the device parameter, which is almost the same as using numpy\n",
    "\n",
    "Note: For closed triangular meshes of genus 0, F = 2(V-2) according to the Euler‐Poincaré formula"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0282f9c3-0ae6-4ba5-b9aa-24d6fbf73112",
   "metadata": {},
   "outputs": [],
   "source": [
    "import trimesh\n",
    "\n",
    "if IN_COLAB:\n",
    "    mesh = trimesh.load_remote(\"https://raw.githubusercontent.com/le-Greg/mesh2sh/master/data/spot_color.obj\",\n",
    "                               maintain_order=True, process=False)\n",
    "else:\n",
    "    mesh = trimesh.load_mesh(\"data/spot_color.obj\", \"obj\", maintain_order=True, process=False)\n",
    "\n",
    "mesh.visual = mesh.visual.to_color()  # Convert UV mapping to vertex-wise colors\n",
    "\n",
    "assert mesh.is_watertight  # Watertight (or a closed manifold mesh)\n",
    "assert mesh.euler_number == 2  # Genus 0 (or topology of a sphere)\n",
    "\n",
    "device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\n",
    "verts = torch.from_numpy(mesh.vertices).to(device=device).float()\n",
    "faces = torch.from_numpy(mesh.faces).to(device=device).long()\n",
    "colors = torch.from_numpy(mesh.visual.vertex_colors).to(device=device)\n",
    "colors = colors[..., :3].float() / 255."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23e82737",
   "metadata": {},
   "source": [
    "Pyvista allows us to render meshes in an interactive window. Let's see what our initial mesh looks like. We chose to use colors by vertices instead of UV mapping, which causes some color distortion."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d3c828a-ab8b-4318-b4d9-0d9997a54d11",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pyvista\n",
    "import numpy as np\n",
    "\n",
    "\n",
    "def show_mesh(vertices, faces, vertices_colors=None):\n",
    "    vertices = vertices.detach().cpu().numpy()\n",
    "    faces = faces.detach().cpu().numpy()\n",
    "    \n",
    "    faces = np.concatenate([np.ones([faces.shape[0], 1], dtype=faces.dtype)*3, faces], axis=1)\n",
    "    mesh = pyvista.PolyData(vertices, faces)\n",
    "    has_colors = vertices_colors is not None\n",
    "    if has_colors:\n",
    "        vertices_colors = vertices_colors.detach().cpu().numpy()\n",
    "        if vertices_colors.dtype == np.float32:\n",
    "            vertices_colors = np.clip(vertices_colors * 255, 0, 255).astype(np.uint8)\n",
    "    if IN_COLAB:\n",
    "        # We use itkwidgets\n",
    "        mesh.rotate_z(180, inplace=True)\n",
    "        mesh.rotate_x(270, inplace=True)\n",
    "        plotter = pyvista.PlotterITK(notebook=True)\n",
    "        plotter.add_mesh(mesh, scalars=vertices_colors)\n",
    "        plotter.show(True)\n",
    "    else:\n",
    "        # We use pythreejs\n",
    "        plotter = pyvista.Plotter(notebook=True)\n",
    "        plotter.add_mesh(mesh, scalars=vertices_colors, rgb=has_colors)\n",
    "        plotter.camera_position = 'xy'\n",
    "        plotter.camera.elevation = 10.\n",
    "        plotter.camera.azimuth = 210.\n",
    "        plotter.show(jupyter_backend='pythreejs', window_size=[500, 500])\n",
    "\n",
    "\n",
    "show_mesh(verts, faces, vertices_colors=colors)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b1994460",
   "metadata": {},
   "source": [
    "## From mesh to spherical harmonics\n",
    "\n",
    "The spherical parametrization is obtained by the Mean Curvature Flow. It is conformal (or angle-preserving), and not authalic (or area-preserving), which means that certain complex parts such as the head are very compressed on the sphere. To reconstruct the mesh, it will be necessary to take this distortion into account.\n",
    "\n",
    "In preparation for what follows, the cmcf function returns the projection on the sphere ('proj') but also the initial curvature density and area distortion vectors ('feat').\n",
    "\n",
    "The execution time depends on the size of the step, but if the mesh does not converge, the step_size must be reduced (and max_iters must be big enough)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44a6ad1b-5b29-476d-aa91-442725797796",
   "metadata": {},
   "outputs": [],
   "source": [
    "from mesh2sh.sphere_parametrization import cmcf\n",
    "from mesh2sh.geometry import spherical_to_cartesian\n",
    "\n",
    "proj, feat = cmcf(verts, faces, max_iters=100, step_size=0.05)\n",
    "\n",
    "\n",
    "sphere_verts = torch.stack([*spherical_to_cartesian(torch.ones([proj.shape[0]], device=proj.device),\n",
    "                                                 proj[:, 0], proj[:, 1])], dim=1)\n",
    "show_mesh(sphere_verts, faces, vertices_colors=colors)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2eba97f6",
   "metadata": {},
   "source": [
    "There are 6 attributes that we want to encode with spherical harmonics: the vertex density to imitate the original mesh density, the 2 geometric properties given by cMCF (curvature density, area distortion), and the 3 RGB colors.\n",
    "\n",
    "The approximation is done by interpolating the points on a uniform grid on the sphere, then adapting the coefficients of the spherical harmonics to come as close as possible to the initial function. For the vertex density, the interpolation is done using a kernel density estimation.\n",
    "\n",
    "\n",
    "sh_coefficients is an array of shape \\[6, lmax+1, lmax+1\\], lmax being the degree of harmonics."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e9d2c08-e1d9-4f4c-8b43-bcaf82ca78c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from mesh2sh.spherical_harmonics import approximate_sh_coefficients, approximate_sh_density\n",
    "\n",
    "lmax = 300\n",
    "feature_vector = torch.cat([feat, colors], dim=1)\n",
    "\n",
    "feature_sph_harm = approximate_sh_coefficients(values=feature_vector, theta=proj[:, 0],\n",
    "                                               phi=proj[:, 1], lmax=lmax)\n",
    "density_sph_harm = approximate_sh_density(theta=proj[:, 0], phi=proj[:, 1], lmax=lmax).unsqueeze(0)\n",
    "\n",
    "sh_coefficients = torch.cat([density_sph_harm, feature_sph_harm], dim=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f1bb490",
   "metadata": {},
   "source": [
    "That's it, we have the 2D representation of the surface, which fully describes the initial properties of the model. We can display them as images and see how it looks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "027ec9de-6238-4c5f-8902-638d05aa62c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "plt.figure(figsize=(12, 6))\n",
    "plt.suptitle('Spherical harmonics coefficients amplitude as an image for each attribute')\n",
    "figname = ['Vertex density', 'Mean curvature density', 'Area Distortion', 'Red', 'Green', 'Blue']\n",
    "for dim in range(sh_coefficients.shape[0]):\n",
    "    plt.subplot(231+dim)\n",
    "    plt.imshow(torch.log(sh_coefficients[dim].abs()).cpu().numpy())\n",
    "    plt.title(figname[dim])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ba0ac53",
   "metadata": {},
   "source": [
    "## From spherical harmonics to mesh\n",
    "\n",
    "Before proceeding with the reconstruction, a new mesh must be made. \n",
    "- Rejection sampling sample a certain number of vertices according to the initial density\n",
    "- Centroidal Voronoi Tesselation optimizes the position of the vertices and creates the faces to have a clean isotropic mesh\n",
    "\n",
    "We choose here the number of vertices to be about 10 times higher than the original number of vertices, to be able to represent all the details"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ebe1f9cd-56d0-4a4f-8c52-00ffbf678867",
   "metadata": {},
   "outputs": [],
   "source": [
    "from mesh2sh.centroidal_voronoi_tessellation import cvt\n",
    "from mesh2sh.rejection_sampling import rejection_sphere_sampling\n",
    "from mesh2sh.spherical_harmonics import pdf_function_from_sh_coeffs\n",
    "\n",
    "n_verts_reconstruction = 20000\n",
    "\n",
    "sh_density = sh_coefficients[0]\n",
    "pdf = pdf_function_from_sh_coeffs(sh_density)\n",
    "raw_estimated_verts = rejection_sphere_sampling(pdf, n=n_verts_reconstruction, device=device)\n",
    "est_sph_verts, est_faces = cvt(raw_estimated_verts, pdf, n_iters=5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54cdf16e",
   "metadata": {},
   "source": [
    "Once the positions of the vertices have been defined, the attributes at these positions are sampled from the spherical harmonic coefficients. The shape is reconstructed using reverse cMCF."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "888d288b-ea5a-488d-a681-03a9c1d193e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from mesh2sh.sphere_parametrization import reverse_cmcf\n",
    "from mesh2sh.spherical_harmonics import sample_sh_values\n",
    "\n",
    "\n",
    "estimated_feature_vector = sample_sh_values(sh_coefficients[1:], theta=est_sph_verts[:, 0], phi=est_sph_verts[:, 1])\n",
    "est_curv = estimated_feature_vector[:, :2]\n",
    "est_colors = estimated_feature_vector[:, 2:]\n",
    "\n",
    "est_verts = reverse_cmcf(sphere_verts=est_sph_verts, features_vector=est_curv, faces=est_faces,\n",
    "                         n_iters=100, step_size=0.0005)\n",
    "\n",
    "show_mesh(est_verts, est_faces, vertices_colors=est_colors)"
   ]
  }
 ],
 "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}