visualize gmm.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pylab import *\n",
    "import pymesh\n",
    "from mixture import GaussianMixture\n",
    "\n",
    "def get_centroids(mesh):\n",
    "    # obtain a vertex for each face index\n",
    "    face_vert = mesh.vertices[mesh.faces.reshape(-1),:].reshape((mesh.faces.shape[0],3,-1))\n",
    "    # face_vert is size (faces,3(one for each vert), 3(one for each dimension))\n",
    "    centroids = face_vert.sum(1)/3.0\n",
    "    #face_vert = ((face_vert.shape[0]/SAMPLE_PTS)*(face_vert.reshape((-1,9))-np.repeat(centroids,3,axis=1)) + np.repeat(centroids,3,axis=1)).reshape((-1,3,3))\n",
    "    ABAC = face_vert[:,1:3,:] - face_vert[:,0:1,:]\n",
    "    areas = np.linalg.norm(np.cross(ABAC[:,0,:],ABAC[:,1,:]),axis=1)/2.0\n",
    "    return centroids, areas\n",
    "\n",
    "def get_tri_covar(tris):\n",
    "    covars = []\n",
    "    for face in face_vert:\n",
    "        A = face[0][:,None]\n",
    "        B = face[1][:,None]\n",
    "        C = face[2][:,None]\n",
    "        M = (A+B+C)/3\n",
    "        covars.append(A @ A.T + B @ B.T + C @ C.T - 3* M @ M.T)\n",
    "    return np.array(covars)*(1/12.0)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import transforms3d\n",
    "mesh0 = pymesh.load_mesh(\"bunny/bun_zipper_1000_1.ply\")\n",
    "K = 100\n",
    "com,a = get_centroids(mesh0)\n",
    "center = mesh0.vertices.mean(0)\n",
    "com -= center\n",
    "print(com.shape)\n",
    "face_vert = mesh0.vertices[mesh0.faces.reshape(-1),:].reshape((mesh0.faces.shape[0],3,-1)) - center\n",
    "R = np.identity(3)\n",
    "R = transforms3d.euler.euler2mat(0,np.pi/2,np.pi/2, 'sxyz')\n",
    "R = transforms3d.euler.euler2mat(np.pi,0,np.pi, 'sxyz') @ R\n",
    "\n",
    "face_vert = face_vert @ R\n",
    "com = com @ R\n",
    "\n",
    "data_covar = get_tri_covar(face_vert)\n",
    "print(data_covar.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gm_mesh_kmeans = GaussianMixture(K,init_params='kmeans',tol=1e-9,max_iter=250); gm_mesh_kmeans.set_covars(data_covar); gm_mesh_kmeans.set_areas(a); gm_mesh_kmeans.fit(com); gm_mesh_kmeans.set_covars(None); gm_mesh_kmeans.set_areas(None)\n",
    "#gm_mesh_kmeans = GaussianMixture(K,init_params='kmeans',tol=1e-9,max_iter=250);  gm_mesh_kmeans.fit(com);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib notebook\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "from mpl_toolkits.mplot3d.art3d import Poly3DCollection\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.colors import LightSource\n",
    "\n",
    "\n",
    "fig = plt.figure(figsize=(15,6))\n",
    "ax = fig.add_subplot(1, 2, 1, projection='3d')\n",
    "b2 = np.array(face_vert)\n",
    "minb = b2.reshape((-1,3)).min(0)\n",
    "maxb = b2.reshape((-1,3)).max(0)\n",
    "meanb = b2.reshape((-1,3)).mean(0)\n",
    "span = maxb-minb\n",
    "w = np.max(span) * np.ones(3)\n",
    "pc = Poly3DCollection(face_vert, alpha = 0.75, facecolor='#1f77b4', linewidths=1,edgecolors=np.array([0.7,0.0,0.5,0.4]))\n",
    "ax.add_collection3d(pc)\n",
    "ax.set_xlim(minb[0]-w[0],maxb[0]+w[0])\n",
    "ax.set_ylim(minb[1]-w[1],maxb[1]+w[1])\n",
    "ax.set_zlim(minb[2]-w[2],maxb[2]+w[2])\n",
    "ax.set_proj_type('persp')\n",
    "\n",
    "\n",
    "# Hide grid lines\n",
    "ax.grid(False)\n",
    "\n",
    "\n",
    "\n",
    "# Hide axes ticks\n",
    "ax.set_xticks([])\n",
    "ax.set_yticks([])\n",
    "ax.set_zticks([])\n",
    "ax.view_init(26,-28)\n",
    "plt.axis('off')\n",
    "ax.dist = 3\n",
    "\n",
    "xlm=ax.get_xlim3d() #These are two tupples\n",
    "ylm=ax.get_ylim3d() #we use them in the next\n",
    "zlm=ax.get_zlim3d() #graph to reproduce the magnification from mousing\n",
    "\n",
    "ax.set_xlim3d(xlm[0]-0.04,xlm[1]-0.04)     #Reproduce magnification\n",
    "ax.set_ylim3d(ylm[0],ylm[1])     #...\n",
    "ax.set_zlim3d(zlm[0],zlm[1])     #...\n",
    "\n",
    "ax = fig.add_subplot(1, 2, 2, projection='3d')\n",
    "#plot_gmm_contour(ax, gm_mesh_kmeans.means_, gm_mesh_kmeans.covariances_, gm_mesh_kmeans.weights_)\n",
    "points=10000\n",
    "z = []\n",
    "c = []\n",
    "colors = plt.cm.tab10(np.arange(10)/10) #plt.rcParams['axes.prop_cycle'].by_key()['color']\n",
    "means,covars,weights = com,data_covar,a/np.sum(a)#gm_mesh_kmeans.means_, gm_mesh_kmeans.covariances_, gm_mesh_kmeans.weights_\n",
    "for e,i in enumerate([14,428,412,446,451,411,439,477]):#range(means.shape[0]):\n",
    "    if np.random.rand() < 5:\n",
    "        samples = int(round(points*weights[i]))\n",
    "        z.append(np.random.multivariate_normal(means[i],covars[i],samples))\n",
    "        c += [colors[e%10] for _ in range(samples)]\n",
    "z = np.vstack(z)\n",
    "c = np.vstack(c)\n",
    "\n",
    "ax.set_xticks([])\n",
    "ax.set_yticks([])\n",
    "ax.set_zticks([])\n",
    "ax.view_init(26,-28)\n",
    "plt.axis('off')\n",
    "ax.set_proj_type('persp')\n",
    "\n",
    "ax.dist = 3\n",
    "pc = Poly3DCollection(face_vert, alpha = 0.1, facecolor='#ffffff', linewidths=1,edgecolors=np.array([0.7,0.0,0.5,0.1]))\n",
    "ax.add_collection3d(pc)\n",
    "#ax.scatter(z[:,0],z[:,1],z[:,2],c=c,depthshade=False)\n",
    "\n",
    "ax.set_xlim3d(xlm[0]-0.04,xlm[1]-0.04)     #Reproduce magnification\n",
    "ax.set_ylim3d(ylm[0],ylm[1])     #...\n",
    "ax.set_zlim3d(zlm[0],zlm[1])     #...\n",
    "\n",
    "faces_of_interest = [14,428,412,446,451,411,439,477]\n",
    "for e,i in enumerate(faces_of_interest):#range(means.shape[0]):\n",
    "    mean,covar = means[i],covars[i]\n",
    "    u,s,vt = np.linalg.svd(covar)\n",
    "    coefs = (.002, .002, .002)  # Coefficients in a0/c x**2 + a1/c y**2 + a2/c z**2 = 1 \n",
    "    # Radii corresponding to the coefficients:\n",
    "    rx, ry, rz = 1.5*np.sqrt(s)#s#1/np.sqrt(coefs)\n",
    "    \n",
    "    R_reg = vt.T @ np.diag([1,1,np.linalg.det(vt.T @ u.T)]) @ u.T\n",
    "    \n",
    "    #print(eigs)\n",
    "    # Set of all spherical angles:\n",
    "    u = np.linspace(0, 2 * np.pi, 25)\n",
    "    v = np.linspace(0, np.pi, 25)\n",
    "\n",
    "    # Cartesian coordinates that correspond to the spherical angles:\n",
    "    # (this is the equation of an ellipsoid):\n",
    "    x = rx * np.outer(np.cos(u), np.sin(v)) #+ mean[0]\n",
    "    y = ry * np.outer(np.sin(u), np.sin(v)) #+ mean[1]\n",
    "    z = rz * np.outer(np.ones_like(u), np.cos(v)) #+ mean[2]\n",
    "    \n",
    "    for i in range(len(x)):\n",
    "        for j in range(len(x)):\n",
    "            x[i,j],y[i,j],z[i,j] = np.dot([x[i,j],y[i,j],z[i,j]], vt) + mean    \n",
    "    # Plot:\n",
    "    ax.plot_surface(x,y,z,  color=colors[e%10],shade=False,linewidth=0.0)\n",
    "pc = Poly3DCollection(face_vert[faces_of_interest], alpha = 0.0, facecolor='#ffffff', linewidths=2,edgecolors=np.array([0.7,0.0,0.5,0.5]))\n",
    "ax.add_collection3d(pc)\n",
    "plt.tight_layout()\n",
    "plt.savefig('bunny_head.pdf')\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sz = 4\n",
    "fig = plt.figure(figsize=(1.4475*sz,sz))\n",
    "padw = (1.4475-1)/2\n",
    "#ax = fig.add_subplot(1, 1, 1, projection='3d')#Axes3D(fig)\n",
    "# L B W H\n",
    "ax = Axes3D(fig, [0, -0.15, 1-padw, 1.1], )\n",
    "#plot_gmm_contour(ax, gm_mesh_kmeans.means_, gm_mesh_kmeans.covariances_, gm_mesh_kmeans.weights_)\n",
    "points=10000\n",
    "z = []\n",
    "c = []\n",
    "ax.autoscale_view(tight=True, scalex=False, scaley=False, scalez=False)\n",
    "\n",
    "pc = Poly3DCollection(face_vert, alpha = 0.01, facecolor='#ffffff', linewidths=1,edgecolors=np.array([0.7,0.0,0.5,0.1]))\n",
    "ax.add_collection3d(pc)\n",
    "#ax.scatter(z[:,0],z[:,1],z[:,2],c=c,depthshade=False)\n",
    "\n",
    "means,covars,weights = gm_mesh_kmeans.means_, gm_mesh_kmeans.covariances_, gm_mesh_kmeans.weights_\n",
    "\n",
    "colors = plt.cm.Pastel1(np.arange(means.shape[0])/(means.shape[0])) #plt.rcParams['axes.prop_cycle'].by_key()['color']\n",
    "np.random.shuffle(colors)\n",
    "\n",
    "faces_of_interest = [14,428,412,446,451,411,439,477]\n",
    "for i in range(means.shape[0]):\n",
    "    orig_i = i\n",
    "    mean,covar = means[i],covars[i]\n",
    "    u,s,vt = np.linalg.svd(covar)\n",
    "    coefs = (.002, .002, .002)  # Coefficients in a0/c x**2 + a1/c y**2 + a2/c z**2 = 1 \n",
    "    # Radii corresponding to the coefficients:\n",
    "    rx, ry, rz = 1.7*np.sqrt(s)#s#1/np.sqrt(coefs)\n",
    "    \n",
    "    R_reg = vt.T @ np.diag([1,1,np.linalg.det(vt.T @ u.T)]) @ u.T\n",
    "    \n",
    "    #print(eigs)\n",
    "    # Set of all spherical angles:\n",
    "    u = np.linspace(0, 2 * np.pi, 10)\n",
    "    v = np.linspace(0, np.pi, 10)\n",
    "\n",
    "    # Cartesian coordinates that correspond to the spherical angles:\n",
    "    # (this is the equation of an ellipsoid):\n",
    "    x = rx * np.outer(np.cos(u), np.sin(v)) #+ mean[0]\n",
    "    y = ry * np.outer(np.sin(u), np.sin(v)) #+ mean[1]\n",
    "    z = rz * np.outer(np.ones_like(u), np.cos(v)) #+ mean[2]\n",
    "    \n",
    "    for i in range(len(x)):\n",
    "        for j in range(len(x)):\n",
    "            x[i,j],y[i,j],z[i,j] = np.dot([x[i,j],y[i,j],z[i,j]], vt) + mean    \n",
    "    # Plot:\n",
    "    res = ax.plot_surface(x,y,z,  color=colors[orig_i],shade=True,linewidth=0.0)\n",
    "plt.tight_layout()\n",
    "\n",
    "ax.set_xticks([])\n",
    "ax.set_yticks([])\n",
    "ax.set_zticks([])\n",
    "ax.view_init(39.4,14)\n",
    "plt.axis('off')\n",
    "ax.set_proj_type('persp')\n",
    "ax.dist = 4.2\n",
    "ax.set_xlim(meanb[0]-w[0],meanb[0]+w[0])\n",
    "ax.set_ylim(meanb[1]-w[1],meanb[1]+w[1])\n",
    "ax.set_zlim(meanb[2]-w[2],meanb[2]+w[2])\n",
    "ax.set_proj_type('persp')\n",
    "plt.tight_layout()\n",
    "plt.savefig('bunny_gmm_k100_mesh.pdf')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "res[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}