{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3D Face Plot\n",
    "\n",
    "Attenzione visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The autoreload extension is already loaded. To reload it, use:\n",
      "  %reload_ext autoreload\n"
     ]
    }
   ],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "import os\n",
    "from os.path import join\n",
    "import sys\n",
    "import time\n",
    "import cv2 as cv\n",
    "import numpy as np\n",
    "import imutils\n",
    "import matplotlib.animation\n",
    "%matplotlib notebook\n",
    "from glob import glob\n",
    "from matplotlib import cbook\n",
    "from matplotlib import cm\n",
    "from matplotlib.colors import LightSource\n",
    "from random import randint\n",
    "sys.path.append('/megapixels/3rdparty/face-alignment')\n",
    "import face_alignment\n",
    "import numpy as np\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "import matplotlib.pyplot as plt\n",
    "import mpl_toolkits.mplot3d.axes3d as p3\n",
    "from matplotlib import animation\n",
    "import random\n",
    "from skimage import io\n",
    "from tqdm import tqdm_notebook as tqdm\n",
    "from IPython.display import clear_output\n",
    "from pathlib import Path\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate random hex colors\n",
    "def rhex():\n",
    "  r = lambda: random.randint(0,255)\n",
    "  return '#%02X%02X%02X' % (r(), r(), r())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading the Face Alignment Network(FAN). Please wait...\n",
      "Downloading the Face Alignment depth Network (FAN-D). Please wait...\n"
     ]
    }
   ],
   "source": [
    "# init 3d face\n",
    "# Run the 3D face alignment on a test image, without CUDA.\n",
    "fa = face_alignment.FaceAlignment(face_alignment.LandmarksType._3D, \n",
    "                                  enable_cuda=True, flip_input=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "100\n"
     ]
    }
   ],
   "source": [
    "data_bodega = '/megapixels/data_bodega/'\n",
    "fp = join(data_bodega,'images/senators/*.jpg')\n",
    "face_files = glob(fp,recursive=True)\n",
    "print(len(face_files))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_3d_face(fpath,output_dir):\n",
    "  # load image\n",
    "  im = io.imread(fpath)\n",
    "  \n",
    "  # generate 3d predictions\n",
    "  preds = fa.get_landmarks(im)\n",
    "  if preds is None:\n",
    "    return\n",
    "  preds = preds[-1]\n",
    "  \n",
    "  # plot style\n",
    "  num_frames = 60\n",
    "  lw = 2 # line weight\n",
    "  bg_color = '#%02x%02x%02x' % (60,59,110)\n",
    "  mark_clr = '#%02x%02x%02x' % (255,255,255) \n",
    "  mark_type='$\\star$'\n",
    "  mark_size = 20\n",
    "  print(bg_color)\n",
    "  \n",
    "  dpi = 72\n",
    "  figsize = (16,16)\n",
    "  \n",
    "  # center x,y,z\n",
    "  xmm = (np.min(preds[:,0]),np.max(preds[:,0]))\n",
    "  ymm = (np.min(preds[:,1]),np.max(preds[:,1]))\n",
    "  zmm = (np.min(preds[:,2]),np.max(preds[:,2]))\n",
    "  \n",
    "  preds_orig = preds.copy()\n",
    "  xmm_sc = (1.2*np.min(preds[:,0]),1.2*np.max(preds_orig[:,0]))\n",
    "  xmm = (np.min(preds_orig[:,0]),np.max(preds_orig[:,0]))\n",
    "  ymm = (np.min(preds_orig[:,1]),np.max(preds_orig[:,1]))\n",
    "  zmm = (np.min(preds_orig[:,2]),np.max(preds_orig[:,2]))\n",
    "  \n",
    "  # swap the y and z components to improve 3d rotation angles\n",
    "  preds = np.zeros_like(preds_orig).astype(np.uint8)\n",
    "  for i,p in enumerate(preds_orig):\n",
    "    x,y,z = p\n",
    "    #preds[i] = np.array([x - xmm[0], y - ymm[0], z - zmm[0]]) # ?\n",
    "    preds[i] = np.array([x - xmm[0], z - zmm[0], y - ymm[0]])\n",
    "  \n",
    "  # Create plot\n",
    "  fig = plt.figure(figsize=figsize,dpi=dpi)\n",
    "  fig.tight_layout()\n",
    "  ax = fig.add_subplot(111, projection='3d')\n",
    "  ax.set_facecolor(bg_color) # background color\n",
    "  \n",
    "  preds_plot = np.zeros_like(preds)\n",
    "  for i,p in enumerate(preds):\n",
    "    x,y,z = p\n",
    "    preds_plot[i] = np.array([x,y,z])\n",
    "\n",
    "  \n",
    "  # scatter plot the dots\n",
    "  ax.plot3D(preds_plot[:17,0]*1.2,preds_plot[:17,1], preds_plot[:17,2],\n",
    "            marker=mark_type,markersize=mark_size,color=mark_clr,linewidth=lw)\n",
    "  ax.plot3D(preds_plot[17:22,0]*1.2,preds_plot[17:22,1],preds_plot[17:22,2],\n",
    "            marker=mark_type,markersize=mark_size,color=mark_clr,linewidth=lw)\n",
    "  ax.plot3D(preds_plot[22:27,0]*1.2,preds_plot[22:27,1],preds_plot[22:27,2], \n",
    "            marker=mark_type,markersize=mark_size,color=mark_clr,linewidth=lw)\n",
    "  ax.plot3D(preds_plot[27:31,0]*1.2,preds_plot[27:31,1],preds_plot[27:31,2],\n",
    "            marker=mark_type,markersize=mark_size,color=mark_clr,linewidth=lw)\n",
    "  ax.plot3D(preds_plot[31:36,0]*1.2,preds_plot[31:36,1],preds_plot[31:36,2],\n",
    "            marker=mark_type,markersize=mark_size,color=mark_clr,linewidth=lw)\n",
    "  ax.plot3D(preds_plot[36:42,0]*1.2,preds_plot[36:42,1],preds_plot[36:42,2],\n",
    "            marker=mark_type,markersize=mark_size,color=mark_clr,linewidth=lw)\n",
    "  ax.plot3D(preds_plot[42:48,0]*1.2,preds_plot[42:48,1],preds_plot[42:48,2],\n",
    "            marker=mark_type,markersize=mark_size,color=mark_clr,linewidth=lw)\n",
    "  ax.plot3D(preds_plot[48:,0]*1.2,preds_plot[48:,1],preds_plot[48:,2],\n",
    "            marker=mark_type,markersize=mark_size,color=mark_clr, linewidth=lw)\n",
    "\n",
    "  \n",
    "  rh = '#ffffff'\n",
    "  ax.scatter(preds_plot[:,0]*1.2,preds_plot[:,1],preds_plot[:,2],c=rh, alpha=1.0, s=35, edgecolor=rh)\n",
    "  \n",
    "  # center points\n",
    "  cx = ((xmm[0] - xmm[1]) // 2) + xmm[1]\n",
    "  cy = ((ymm[1] - ymm[0]) // 2) + ymm[0]\n",
    "  cz = ((zmm[1] - zmm[0]) // 2) + zmm[0]\n",
    "\n",
    "  #xpts = [cx,cx]\n",
    "  #ypts = [cy,cy]\n",
    "  #zpts = [zmm[0],zmm[1]]\n",
    "\n",
    "  ax.view_init(elev=120., azim=70.)\n",
    "  \n",
    "  ax.set_xticks([])\n",
    "  ax.set_yticks([])\n",
    "  ax.set_zticks([])\n",
    "  ax.set_axis_off()\n",
    "  ax.set_xlabel('x')\n",
    "  ax.set_ylabel('y')\n",
    "  ax.set_zlabel('z')\n",
    "\n",
    "  # write animation to disk\n",
    "  fname_out = join(output_dir, '{}.gif'.format(Path(fpath).stem))\n",
    "  print(fname_out)\n",
    "  # rotation increments\n",
    "  phi = np.linspace(0, 2*np.pi, num_frames)\n",
    "\n",
    "  def update(phi):\n",
    "      ax.view_init(180,phi*180./np.pi)\n",
    "  \n",
    "  ani = matplotlib.animation.FuncAnimation(fig, update, frames=phi)\n",
    "  ani.save(fname_out, writer='imagemagick', fps=10)\n",
    "  #del ani"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_dir = join(data_bodega,'output','senators_3d_points')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "for f in tqdm(face_files):\n",
    "  generate_3d_face(f,output_dir)\n",
    "  clear_output()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cv.L"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:megapixels]",
   "language": "python",
   "name": "conda-env-megapixels-py"
  },
  "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}