mediapipe_face_mesh_to_blendshapes.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "s_PwLpmCnBYW"
   },
   "source": [
    "\n",
    "Copyright (c) 2023 Haiba Labs\n",
    "\n",
    "Author: James Ritts <james@haibalabs.com>\n",
    "\n",
    "This notebook trains a simple pytorch model to map from [MediaPipe face mesh](http://solutions.mediapipe.dev/face_mesh) landmarks to [ARKit-compatible blendshapes](https://developer.apple.com/documentation/arkit/arfaceanchor/blendshapelocation).\n",
    "\n",
    "### [Click here to open the demo.](https://haibalabs.github.io/face-mesh-to-blendshapes/test/mediapipe_to_arkit.html)\n",
    "\n",
    "### Caveats\n",
    "\n",
    "- We wish to train on object space geo so it doesn't have to learn what a face pose looks like in every possible head orientation. Unfortunately MediaPipe's output is only given in [screen coordinates](https://www.cse.iitd.ac.in/~suban/vision/affine/node5.html). Its mesh is also stretched to conform to the silhouette of the face in the input image. The function normalize_landmarks() tries to undo these effects.\n",
    "- The function convert_landmarks_to_model_input() uses normalize_landmarks in order to convert from raw MediaPipe output to the NN input vector. This function needs to be ported to any environment where the model is run.\n",
    "- MediaPipe isn't able to signal every blendshape. These should be forced to zero at runtime and possibly others as well: jawForward, jawRight, jawLeft, mouthDimpleRight, mouthDimpleLeft, cheekPuff, tongueOut.\n",
    "\n",
    "\n",
    "### Format\n",
    "\n",
    "The order of blendshape values in the model output is:\n",
    "\n",
    "```\n",
    "eyeBlinkRight, eyeLookDownRight, eyeLookInRight, eyeLookOutRight, eyeLookUpRight, eyeSquintRight, eyeWideRight, eyeBlinkLeft, eyeLookDownLeft, eyeLookInLeft, eyeLookOutLeft, eyeLookUpLeft, eyeSquintLeft, eyeWideLeft, jawForward, jawRight, jawLeft, jawOpen, mouthClose, mouthFunnel, mouthPucker, mouthRight, mouthLeft, mouthSmileRight, mouthSmileLeft, mouthFrownRight, mouthFrownLeft, mouthDimpleRight, mouthDimpleLeft, mouthStretchRight, mouthStretchLeft, mouthRollLower, mouthRollUpper, mouthShrugLower, mouthShrugUpper, mouthPressRight, mouthPressLeft, mouthLowerDownRight, mouthLowerDownLeft, mouthUpperUpRight, mouthUpperUpLeft, browDownRight, browDownLeft, browInnerUp, browOuterUpRight, browOuterUpLeft, cheekPuff, cheekSquintRight, cheekSquintLeft, noseSneerRight, noseSneerLeft, tongueOut\n",
    "```\n",
    "\n",
    "Note MediaPipe isn't capable of signaling every blendshape. For example, these should be probably forced to zero at runtime, and possibly others:\n",
    "\n",
    "```\n",
    "jawForward, jawRight, jawLeft, mouthDimpleRight, mouthDimpleLeft, cheekPuff, tongueOut\n",
    "```\n",
    "\n",
    "Training data has this folder structure:\n",
    "- my_first_dataset\n",
    "  - neutral.jpg\n",
    "  - my_first_dataset.csv\n",
    "  - my_first_dataset_000000.jpg\n",
    "  - my_first_dataset_000001.jpg\n",
    "  - my_first_dataset_000002.jpg\n",
    "  - ...\n",
    "- my_second_dataset\n",
    "- sets.txt\n",
    "\n",
    "The file sets.txt should contain the folder names of all training datasets:\n",
    "```\n",
    "my_first_dataset\n",
    "my_second_dataset\n",
    "...\n",
    "```\n",
    "\n",
    "Each dataset must have a calibration photo depicting a neutral facial expression: **neutral.jpg**.  The model is trained on object space offsets from the neutral pose.\n",
    "\n",
    "All images in a set should have approximately the same head transform.\n",
    "\n",
    "Each dataset also has a CSV file containing a header row followed by labels (blendshape values) for each input image:\n",
    "```\n",
    "eyeBlinkRight,eyeLookDownRight,eyeLookInRight,eyeLookOutRight,eyeLookUpRight,eyeSquintRight,eyeWideRight,eyeBlinkLeft,eyeLookDownLeft,eyeLookInLeft,eyeLookOutLeft,eyeLookUpLeft,eyeSquintLeft,eyeWideLeft,jawForward,jawRight,jawLeft,jawOpen,mouthClose,mouthFunnel,mouthPucker,mouthRight,mouthLeft,mouthSmileRight,mouthSmileLeft,mouthFrownRight,mouthFrownLeft,mouthDimpleRight,mouthDimpleLeft,mouthStretchRight,mouthStretchLeft,mouthRollLower,mouthRollUpper,mouthShrugLower,mouthShrugUpper,mouthPressRight,mouthPressLeft,mouthLowerDownRight,mouthLowerDownLeft,mouthUpperUpRight,mouthUpperUpLeft,browDownRight,browDownLeft,browInnerUp,browOuterUpRight,browOuterUpLeft,cheekPuff,cheekSquintRight,cheekSquintLeft,noseSneerRight,noseSneerLeft,tongueOut\n",
    "0.039,0.103,0.044,0.000,0.000,0.000,0.000,0.039,0.104,0.000,0.000,0.000,0.000,0.000,0.000,0.000,0.000,0.000,0.000,0.000,0.000,0.000,0.010,0.010,0.027,0.000,0.000,0.002,0.003,0.000,0.000,0.000,0.000,0.000,0.000,0.000,0.000,0.000,0.000,0.000,0.000,0.015,0.014,0.000,0.000,0.000,0.007,0.000,0.000,0.000,0.000,0.000\n",
    "0.038,0.091,0.049,0.000,0.000,0.000,0.000,0.038,0.092,0.000,0.000,0.000,0.000,0.000,0.000,0.000,0.000,0.000,0.000,0.000,0.000,0.000,0.010,0.011,0.027,0.000,0.000,0.002,0.004,0.000,0.000,0.000,0.000,0.000,0.000,0.000,0.000,0.000,0.000,0.000,0.000,0.014,0.014,0.000,0.000,0.000,0.007,0.000,0.000,0.000,0.000,0.000\n",
    "...\n",
    "```\n",
    "\n",
    "To do:\n",
    "- cull shapes from NN output and training which MP can't signal\n",
    "- cull training examples which don't signal shapes that MP can detect (programmically; already did a rough manual pass)\n",
    "\n",
    "Relevant links:\n",
    "- https://arxiv.org/pdf/2006.10962.pdf\n",
    "- https://developers.googleblog.com/2020/09/mediapipe-3d-face-transform.html\n",
    "- https://github.com/google/mediapipe/tree/master/mediapipe/modules/face_geometry/data\n",
    "- https://github.com/google/mediapipe/issues/2867\n",
    "- https://stackoverflow.com/questions/69858216/mediapipe-facemesh-vertices-mapping\n",
    "- https://github.com/Rassibassi/mediapipeFacegeometryPython/blob/main/face_geometry.py\n",
    "- https://github.com/google/mediapipe/blob/a908d668c730da128dfa8d9f6bd25d519d006692/mediapipe/modules/face_geometry/data/canonical_face_model_uv_visualization.png"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ic0X9m5n8kMP"
   },
   "source": [
    "---\n",
    "**Configuration**\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "id": "mz0FynoN4YBT"
   },
   "outputs": [],
   "source": [
    "import math\n",
    "\n",
    "if 'google.colab' in str(get_ipython()):\n",
    "    IS_LOCAL_ENVIRONMENT = False\n",
    "else:\n",
    "    IS_LOCAL_ENVIRONMENT = True\n",
    "\n",
    "# Root folder with training data\n",
    "# Cache and model files will be written here\n",
    "if IS_LOCAL_ENVIRONMENT:\n",
    "    DATA_PATH = 'Face/Training/'\n",
    "else:\n",
    "    DATA_PATH = '/content/drive/MyDrive/Colab/Face/Training/'\n",
    "\n",
    "# Whether to enable MediaPipe's refineLandmarks feature\n",
    "REFINE_LANDMARKS = True\n",
    "\n",
    "# Regenerate training data from the contents of DATA_PATH\n",
    "RUN_MEDIAPIPE_GEN = True\n",
    "\n",
    "# Train the NN\n",
    "RUN_TRAINING = True # train NN\n",
    "\n",
    "# Run an interactive test of the model input preprocessing\n",
    "RUN_NORMALIZE_TEST = True\n",
    "\n",
    "# Vector sizes\n",
    "NUM_LANDMARKS = 225 # see MP_MOUTH_INDICES, ...\n",
    "NN_INPUT_SIZE = NUM_LANDMARKS * 2\n",
    "NN_OUTPUT_SIZE = 52\n",
    "\n",
    "# Hyperparams\n",
    "NN_HIDDEN_SIZE = math.floor(((2 * NN_INPUT_SIZE) / 3) + NN_OUTPUT_SIZE)\n",
    "NN_EPOCHS = 75\n",
    "\n",
    "# Misc\n",
    "SHOW_IMAGE_HEIGHT = 640\n",
    "SHOW_IMAGE_WIDTH = 640\n",
    "\n",
    "# MediaPipe outputs 478 landmarks, from which we cull all but 225 \"useful\" ones\n",
    "MP_MOUTH_INDICES = [\n",
    "    # lips\n",
    "    0, 11, 12, 13, 14, 15, 16, 17, 37, 38, 39, 40, 41, 42, 61, 62, 72, 73, 74, 76, 77, 78, 80, 81, 82, 84, 85, 86, 87, 88, 89, 90, 91, 95, 96, 146, 178, 179, 180, 181, 183, 184, 185, 191, 267, 268, 269, 270, 271, 272, 291, 292, 302, 303, 304, 306, 307, 308, 310, 311, 312, 314, 315, 316, 317, 318, 319, 320, 321, 324, 325, 375, 402, 403, 404, 405, 407, 408, 409,\n",
    "    # first ring\n",
    "    18, 43, 57, 83, 92, 106, 164, 165, 167, 182, 186, 273, 287, 313, 322, 335, 391, 393, 406, 410\n",
    "]\n",
    "\n",
    "MP_LEFT_EYE_INDICES = [\n",
    "    # eye socket\n",
    "    7, 33, 133, 144, 145, 153, 154, 155, 157, 158, 159, 160, 161, 163, 173, 246,\n",
    "    # first ring around eye\n",
    "    22, 23, 24, 25, 26, 27, 28, 29, 30, 56, 110, 112, 130, 190, 243, 247,\n",
    "    # second ring around eye\n",
    "    31, 113, 189, 221, 222, 223, 224, 225, 226, 228, 229, 230, 231, 232, 233, 244,\n",
    "    # brow row 1\n",
    "    46, 52, 53, 55, 65,\n",
    "    # brow row 2\n",
    "    63, 66, 70, 105, 107,\n",
    "    # brow row 3\n",
    "    68, 69, 71, 104, 108\n",
    "]\n",
    "\n",
    "MP_RIGHT_EYE_INDICES = [\n",
    "    # eye socket\n",
    "    249, 263, 362, 373, 374, 380, 381, 382, 384, 385, 386, 387, 388, 390, 398, 466,\n",
    "    # first ring around eye\n",
    "    252, 253, 254, 255, 256, 257, 258, 259, 260, 286, 339, 341, 359, 414, 463, 467,\n",
    "    # second ring around eye\n",
    "    261, 342, 413, 441, 442, 443, 444, 445, 446, 448, 449, 450, 451, 452, 453, 464,\n",
    "    # brow row 1\n",
    "    276, 282, 283, 285, 295,\n",
    "    # brow row 2\n",
    "    293, 296, 300, 334, 336,\n",
    "    # brow row 3\n",
    "    298, 299, 301, 333, 337\n",
    "]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "bluIfskl4YhW"
   },
   "source": [
    "---\n",
    "**Dependencies**\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "id": "1FEA-_sYzboT"
   },
   "outputs": [],
   "source": [
    "# Scratch local setup\n",
    "#!python --version\n",
    "# !conda create --name myenv\n",
    "# !conda activate colab_env\n",
    "#!pip show torch\n",
    "# !pip show torchvision\n",
    "#!pip install torch==1.13.0+cu116 torchvision==0.14.0+cu116 -f https://download.pytorch.org/whl/torch_stable.html\n",
    "# !pip install numpy==1.21.6\n",
    "# !pip install --extra-index-url https://download.pytorch.org/whl/cu113/ \"torch==1.12.1+cu113\"\n",
    "# !pip install pandas==1.3.5\n",
    "\n",
    "if not IS_LOCAL_ENVIRONMENT:\n",
    "    !pip install mediapipe\n",
    "    !pip install pytorch_lightning\n",
    "    !pip install torchviz\n",
    "    !pip install torchvision\n",
    "    !pip install matplotlib\n",
    "    !git clone https://github.com/Rassibassi/mediapipeDemos\n",
    "\n",
    "import pytorch_lightning as pl\n",
    "from torchmetrics import Accuracy\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torchvision\n",
    "from torchvision import transforms\n",
    "from torch.utils.data import DataLoader\n",
    "from torch.utils.data import random_split\n",
    "from torch.utils.data.dataset import Dataset\n",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "import cv2\n",
    "import numpy as np\n",
    "import mediapipe as mp\n",
    "import os\n",
    "import glob\n",
    "import pandas as pd\n",
    "import copy\n",
    "import pickle\n",
    "from ast import Yield\n",
    "from IPython.display import display, Javascript\n",
    "from base64 import b64decode\n",
    "from torchviz import make_dot\n",
    "\n",
    "mp_face_mesh = mp.solutions.face_mesh\n",
    "mp_face_mesh_connections = mp.solutions.face_mesh_connections\n",
    "mp_drawing = mp.solutions.drawing_utils \n",
    "mp_drawing_styles = mp.solutions.drawing_styles\n",
    "drawing_spec = mp_drawing.DrawingSpec(thickness=1, circle_radius=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WGl55h4g2BWL"
   },
   "source": [
    "---\n",
    "**Utilities**\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "id": "jgceGo3g2DmP"
   },
   "outputs": [],
   "source": [
    "if not IS_LOCAL_ENVIRONMENT:\n",
    "    from google.colab.output import eval_js\n",
    "    from google.colab.patches import cv2_imshow\n",
    "\n",
    "# https://graphics.pixar.com/library/OrthonormalB/paper.pdf\n",
    "def revisedONB(n):\n",
    "    if n[2] < 0:\n",
    "        a = 1.0 / (1.0 - n[2])\n",
    "        b = n[0] * n[1] * a\n",
    "        return [[1.0 - n[0] * n[0] * a, -b, n[0]], [b, n[1] * n[1]*a - 1.0, -n[1]]]\n",
    "    else:\n",
    "        a = 1.0 / (1.0 + n[2])\n",
    "        b = -n[0] * n[1] * a\n",
    "        return [[1.0 - n[0] * n[0] * a, b, -n[0]], [b, 1.0 - n[1] * n[1] * a, -n[1]]]\n",
    "\n",
    "# Get a photo using the webcam\n",
    "def take_photo(filename='photo.jpg', quality=0.8):\n",
    "    js = Javascript('''\n",
    "        async function takePhoto(quality) {\n",
    "            const div = document.createElement('div');\n",
    "            const capture = document.createElement('button');\n",
    "            capture.textContent = 'Capture';\n",
    "            div.appendChild(capture);\n",
    "\n",
    "            const video = document.createElement('video');\n",
    "            video.style.display = 'block';\n",
    "            const stream = await navigator.mediaDevices.getUserMedia({video: true});\n",
    "\n",
    "            document.body.appendChild(div);\n",
    "            div.appendChild(video);\n",
    "            video.srcObject = stream;\n",
    "            await video.play();\n",
    "\n",
    "            // Resize the output to fit the video element.\n",
    "            google.colab.output.setIframeHeight(document.documentElement.scrollHeight, true);\n",
    "\n",
    "            // Wait for Capture to be clicked.\n",
    "            await new Promise((resolve) => capture.onclick = resolve);\n",
    "\n",
    "            const canvas = document.createElement('canvas');\n",
    "            canvas.width = video.videoWidth;\n",
    "            canvas.height = video.videoHeight;\n",
    "            canvas.getContext('2d').drawImage(video, 0, 0);\n",
    "            stream.getVideoTracks()[0].stop();\n",
    "            div.remove();\n",
    "            return canvas.toDataURL('image/jpeg', quality);\n",
    "        }\n",
    "        ''')\n",
    "    display(js)\n",
    "    data = eval_js('takePhoto({})'.format(quality))\n",
    "    binary = b64decode(data.split(',')[1])\n",
    "    with open(filename, 'wb') as f:\n",
    "        f.write(binary)\n",
    "    return filename\n",
    "\n",
    "# Scale and display an image\n",
    "def show_image(image):\n",
    "    h, w = image.shape[:2]\n",
    "    if h < w:\n",
    "        img = cv2.resize(image, (SHOW_IMAGE_WIDTH, math.floor(h/(w/SHOW_IMAGE_WIDTH))))\n",
    "    else:\n",
    "        img = cv2.resize(image, (math.floor(w/(h/SHOW_IMAGE_HEIGHT)), SHOW_IMAGE_HEIGHT))\n",
    "    cv2_imshow(img)\n",
    "\n",
    "# Scale and display an image overlaid with mediapipe landbarks\n",
    "def show_landmarks(image, multi_face_landmarks):\n",
    "    copy = image.copy()\n",
    "    for face_landmarks in multi_face_landmarks:\n",
    "        mp_drawing.draw_landmarks(\n",
    "                image=copy,\n",
    "                landmark_list=face_landmarks,\n",
    "                connections=mp_face_mesh.FACEMESH_TESSELATION,\n",
    "                landmark_drawing_spec=None,\n",
    "                connection_drawing_spec=mp_drawing_styles\n",
    "                .get_default_face_mesh_tesselation_style())\n",
    "        mp_drawing.draw_landmarks(\n",
    "                image=copy,\n",
    "                landmark_list=face_landmarks,\n",
    "                connections=mp_face_mesh.FACEMESH_CONTOURS,\n",
    "                landmark_drawing_spec=None,\n",
    "                connection_drawing_spec=mp_drawing_styles\n",
    "                .get_default_face_mesh_contours_style())\n",
    "        if REFINE_LANDMARKS:\n",
    "            mp_drawing.draw_landmarks(\n",
    "                    image=copy,\n",
    "                    landmark_list=face_landmarks,\n",
    "                    connections=mp_face_mesh.FACEMESH_IRISES,\n",
    "                    landmark_drawing_spec=None,\n",
    "                    connection_drawing_spec=mp_drawing_styles\n",
    "                    .get_default_face_mesh_iris_connections_style())\n",
    "    show_image(copy)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "52ZYfSJQ_7D_"
   },
   "source": [
    "---\n",
    "**Preprocessing**\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "id": "8F3DgcGw_ylz"
   },
   "outputs": [
    {
     "data": {
      "application/javascript": "\n        async function takePhoto(quality) {\n            const div = document.createElement('div');\n            const capture = document.createElement('button');\n            capture.textContent = 'Capture';\n            div.appendChild(capture);\n\n            const video = document.createElement('video');\n            video.style.display = 'block';\n            const stream = await navigator.mediaDevices.getUserMedia({video: true});\n\n            document.body.appendChild(div);\n            div.appendChild(video);\n            video.srcObject = stream;\n            await video.play();\n\n            // Resize the output to fit the video element.\n            google.colab.output.setIframeHeight(document.documentElement.scrollHeight, true);\n\n            // Wait for Capture to be clicked.\n            await new Promise((resolve) => capture.onclick = resolve);\n\n            const canvas = document.createElement('canvas');\n            canvas.width = video.videoWidth;\n            canvas.height = video.videoHeight;\n            canvas.getContext('2d').drawImage(video, 0, 0);\n            stream.getVideoTracks()[0].stop();\n            div.remove();\n            return canvas.toDataURL('image/jpeg', quality);\n        }\n        ",
      "text/plain": [
       "<IPython.core.display.Javascript object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "name 'eval_js' is not defined\n"
     ]
    }
   ],
   "source": [
    "# MediaPipe's face mesh indices:\n",
    "#   https://github.com/google/mediapipe/blob/master/mediapipe/modules/face_geometry/data/canonical_face_model_uv_visualization.png\n",
    "# \n",
    "# Landmark coordinate system is +X to the right, +Y down, and +Z pointing into the screen.\n",
    "# \n",
    "#            _\n",
    "#            /| +Z\n",
    "#           /\n",
    "#          /\n",
    "#         o-----> +X\n",
    "#         |\n",
    "#         |\n",
    "#         \\/ +Y\n",
    "\n",
    "def screen_align_and_normalize(arr, xf):\n",
    "    # screen align\n",
    "    pivot = np.mean(arr, axis=0)\n",
    "    arr = [p - pivot for p in arr]\n",
    "    arr = [xf.dot(p).tolist()[0] for p in arr]\n",
    "    arr = [p + pivot for p in arr]\n",
    "    # fill unit cube\n",
    "    amin = np.amin(arr, axis=0)\n",
    "    amax = np.amax(arr, axis=0)\n",
    "    for p in arr:\n",
    "        p[0] = (p[0] - amin[0]) / (amax[0] - amin[0])\n",
    "        p[1] = (p[1] - amin[1]) / (amax[1] - amin[1])\n",
    "        p[2] = (p[2] - amin[2]) / (amax[2] - amin[2])\n",
    "    return arr\n",
    "\n",
    "# Get world=>object map for the full mesh using a basis calculated from forehead points\n",
    "def calc_forehead_xf(arr):\n",
    "    # average a few normals at the top of the head to get a forward vec\n",
    "    v0 = np.subtract(arr[151], arr[10])\n",
    "    v1 = np.subtract(arr[338], arr[10])\n",
    "    v2 = np.subtract(arr[109], arr[10])\n",
    "    f1 = np.cross(v1, v0)\n",
    "    f2 = np.cross(v0, v2)\n",
    "    f1 = f1 / np.linalg.norm(f1)\n",
    "    f2 = f2 / np.linalg.norm(f2)\n",
    "    vfw = 0.5 * (f1 + f2)\n",
    "    # up\n",
    "    vup = np.subtract(arr[151], arr[10])\n",
    "    vup = vup / np.linalg.norm(vup)\n",
    "    # right\n",
    "    vrt = np.cross(vup, vfw)\n",
    "    vup = np.cross(vfw, vrt)\n",
    "    return np.matrix([vrt, vup, vfw]) # 3x3\n",
    "\n",
    "# Get world=>object map for the left eye\n",
    "def calc_simple_xf(arr, idxF0, idxF1, idxF2, idxR0, idxR1):\n",
    "    # fwd\n",
    "    v0 = np.subtract(arr[idxF2], arr[idxF0])\n",
    "    v1 = np.subtract(arr[idxF1], arr[idxF0])\n",
    "    vfw = np.cross(v1, v0)\n",
    "    vfw = vfw / np.linalg.norm(vfw)\n",
    "    # right\n",
    "    vrt = np.subtract(arr[idxR1], arr[idxR0])\n",
    "    vrt = vrt / np.linalg.norm(vrt)\n",
    "    # up\n",
    "    vup = np.cross(vfw, vrt)\n",
    "    vrt = np.cross(vup, vfw)\n",
    "    return np.matrix([vrt, vup, vfw]) # 3x3\n",
    "\n",
    "# Get world=>object map for the left eye\n",
    "def calc_left_eye_xf(arr):\n",
    "    return calc_simple_xf(arr, 23, 22, 230, 33, 133)\n",
    "\n",
    "# Get world=>object map for the left eye\n",
    "def calc_right_eye_xf(arr):\n",
    "    return calc_simple_xf(arr, 253, 450, 252, 362, 263)\n",
    "\n",
    "# Get world=>object map for the left eye\n",
    "def calc_mouth_xf(arr):\n",
    "    # average a few normals at the top of the head to get a forward vec\n",
    "    v0 = np.subtract(arr[164], arr[2])\n",
    "    v1 = np.subtract(arr[326], arr[2])\n",
    "    v2 = np.subtract(arr[97], arr[2])\n",
    "    f1 = np.cross(v1, v0)\n",
    "    f2 = np.cross(v0, v2)\n",
    "    f1 = f1 / np.linalg.norm(f1)\n",
    "    f2 = f2 / np.linalg.norm(f2)\n",
    "    vfw = 0.5 * (f1 + f2)\n",
    "    # right\n",
    "    vrt = np.subtract(arr[312], arr[82])\n",
    "    vrt = vrt / np.linalg.norm(vrt)\n",
    "    # up\n",
    "    vup = np.cross(vfw, vrt)\n",
    "    vrt = np.cross(vup, vfw)\n",
    "    return np.matrix([vrt, vup, vfw]) # 3x3\n",
    "\n",
    "def normalize_landmarks(multi_face_landmarks):\n",
    "    landmarks = multi_face_landmarks[0].landmark\n",
    "    arr = [[l.x, l.y, l.z] for l in landmarks]\n",
    "\n",
    "    arrMouth = [arr[idx] for idx in MP_MOUTH_INDICES]\n",
    "    arrLeftEye = [arr[idx] for idx in MP_LEFT_EYE_INDICES]\n",
    "    arrRightEye = [arr[idx] for idx in MP_RIGHT_EYE_INDICES]\n",
    "\n",
    "    arrMouth = screen_align_and_normalize(arrMouth, calc_mouth_xf(arr))\n",
    "    arrLeftEye = screen_align_and_normalize(arrLeftEye, calc_left_eye_xf(arr))\n",
    "    arrRightEye = screen_align_and_normalize(arrRightEye, calc_right_eye_xf(arr))\n",
    "\n",
    "    for idx in range(len(landmarks)):\n",
    "        landmarks[idx].x = 0\n",
    "        landmarks[idx].y = 0\n",
    "        landmarks[idx].z = 0\n",
    "\n",
    "    src = 0\n",
    "    for idx in MP_MOUTH_INDICES:\n",
    "        landmarks[idx].x = arrMouth[src][0]\n",
    "        landmarks[idx].y = arrMouth[src][1]\n",
    "        landmarks[idx].z = arrMouth[src][2]\n",
    "        src = src + 1\n",
    "\n",
    "    src = 0\n",
    "    for idx in MP_LEFT_EYE_INDICES:\n",
    "        landmarks[idx].x = arrLeftEye[src][0]\n",
    "        landmarks[idx].y = arrLeftEye[src][1]\n",
    "        landmarks[idx].z = arrLeftEye[src][2]\n",
    "        src = src + 1\n",
    "\n",
    "    src = 0\n",
    "    for idx in MP_RIGHT_EYE_INDICES:\n",
    "        landmarks[idx].x = arrRightEye[src][0]\n",
    "        landmarks[idx].y = arrRightEye[src][1]\n",
    "        landmarks[idx].z = arrRightEye[src][2]\n",
    "        src = src + 1\n",
    "\n",
    "    arr = arrMouth + arrLeftEye + arrRightEye;\n",
    "    arr = [[i[0], i[1]] for i in arr] # Z -> 0\n",
    "    return arr\n",
    "\n",
    "def get_normalized_landmarks(path, show=False):\n",
    "    image = cv2.imread(path)\n",
    "    with mp_face_mesh.FaceMesh(\n",
    "            static_image_mode=True,\n",
    "            refine_landmarks=REFINE_LANDMARKS,\n",
    "            max_num_faces=1,\n",
    "            min_detection_confidence=0.5) as face_mesh:\n",
    "        results = face_mesh.process(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))\n",
    "        arr = normalize_landmarks(results.multi_face_landmarks)\n",
    "        if (show):\n",
    "            show_landmarks(image, results.multi_face_landmarks)\n",
    "        return arr\n",
    "\n",
    "def convert_landmarks_to_model_input(multi_face_landmarks, neutral_normalized_landmarks):\n",
    "    arr = normalize_landmarks(multi_face_landmarks)\n",
    "    return np.subtract(arr, neutral_normalized_landmarks)\n",
    "\n",
    "# Test code\n",
    "if RUN_NORMALIZE_TEST:\n",
    "    from IPython.display import Image\n",
    "    try:\n",
    "        #test_path = '/content/drive/MyDrive/Colab/Face/Training/eyes1/eyes1_0013.jpg'\n",
    "        test_path = take_photo()\n",
    "        test_image = cv2.imread(test_path)\n",
    "        with mp_face_mesh.FaceMesh(\n",
    "                static_image_mode=True,\n",
    "                refine_landmarks=REFINE_LANDMARKS,\n",
    "                max_num_faces=1,\n",
    "                min_detection_confidence=0.5) as face_mesh:\n",
    "\n",
    "            results = face_mesh.process(cv2.cvtColor(test_image, cv2.COLOR_BGR2RGB))\n",
    "            normalize_landmarks(results.multi_face_landmarks)\n",
    "\n",
    "            show_landmarks(test_image, results.multi_face_landmarks)\n",
    "\n",
    "    except Exception as err:\n",
    "        # Errors will be thrown if the user does not have a webcam or if they do not\n",
    "        # grant the page permission to access it.\n",
    "        print(str(err))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-npL2XwktpaM"
   },
   "source": [
    "---\n",
    "**Model definition**\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "id": "e1Y0KD2wtpCL"
   },
   "outputs": [],
   "source": [
    "class FullyConnectedModel(pl.LightningModule):\n",
    "    def __init__(self, input_size=NN_INPUT_SIZE, output_size=NN_OUTPUT_SIZE, hidden_units=[NN_HIDDEN_SIZE]):\n",
    "        super().__init__()\n",
    "        \n",
    "        all_layers = [nn.Flatten()]\n",
    "        for hidden_unit in hidden_units:\n",
    "            layer = nn.Linear(input_size, hidden_unit)\n",
    "            all_layers.append(layer)\n",
    "            all_layers.append(nn.ReLU()) \n",
    "            input_size = hidden_unit \n",
    " \n",
    "        all_layers.append(nn.Linear(input_size, output_size))\n",
    "        self.writer = SummaryWriter()\n",
    "        self.model = nn.Sequential(*all_layers)\n",
    "        self.epoch = 0\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.model(x)\n",
    "\n",
    "    def training_step(self, batch, batch_idx):\n",
    "        x, y = batch\n",
    "        preds = self(x)\n",
    "        loss = nn.functional.mse_loss(self(x), y)\n",
    "        self.log(\"train_loss\", loss, prog_bar=True)\n",
    "\n",
    "        self.epoch += 1\n",
    "        self.writer.add_scalar(\"Loss/train\", loss, self.epoch)\n",
    "\n",
    "        return loss\n",
    "\n",
    "    def validation_step(self, batch, batch_idx):\n",
    "        x, y = batch\n",
    "        preds = self(x)\n",
    "        loss = nn.functional.mse_loss(self(x), y)\n",
    "        return loss\n",
    "    \n",
    "    def test_step(self, batch, batch_idx):\n",
    "        x, y = batch\n",
    "        preds = self(x)\n",
    "        loss = nn.functional.mse_loss(self(x), y)\n",
    "        return loss\n",
    "\n",
    "    def configure_optimizers(self):\n",
    "        optimizer = torch.optim.Adam(self.parameters(), lr=0.001)\n",
    "        return optimizer\n",
    "\n",
    "class ARKitDataset(Dataset):\n",
    "    def __init__(self, sets):\n",
    "        self.inputs = [torch.from_numpy(np.array(i).flatten()).float() for s in sets for i in s['input']]\n",
    "        self.labels = [torch.from_numpy(np.array(i)).float() for s in sets for i in s['output']]\n",
    "    def __len__(self):\n",
    "        return len(self.inputs)\n",
    "    def __getitem__(self, idx):\n",
    "        return self.inputs[idx], self.labels[idx]\n",
    "\n",
    "class ARKitDataModule(pl.LightningDataModule):\n",
    "    def __init__(self, sets, batch_size = 64, validation_size = 0.1):\n",
    "        super().__init__()\n",
    "        self.sets = sets\n",
    "        self.batch_size = batch_size\n",
    "        self.validation_size = validation_size\n",
    "\n",
    "    def prepare_data(self):\n",
    "        self.data = ARKitDataset(self.sets)\n",
    "        self.data_size = len(self.data)\n",
    "        self.val_size = math.floor(self.validation_size * self.data_size)\n",
    "        self.train_size = self.data_size - self.val_size\n",
    "\n",
    "    def setup(self, stage=None):\n",
    "        if stage == \"fit\" or stage is None:\n",
    "            self.train, self.val = random_split(self.data, [self.train_size, self.val_size])\n",
    "        if stage == \"test\" or stage is None:\n",
    "            self.test = self.data\n",
    "        if stage == \"predict\" or stage is None:\n",
    "            self.predict = self.data\n",
    "\n",
    "    def train_dataloader(self):\n",
    "        return DataLoader(self.train, batch_size=self.batch_size)\n",
    "\n",
    "    def val_dataloader(self):\n",
    "        return DataLoader(self.val, batch_size=self.batch_size)\n",
    "\n",
    "    def test_dataloader(self):\n",
    "        return DataLoader(self.test, batch_size=self.batch_size)\n",
    "\n",
    "    def predict_dataloader(self):\n",
    "        return DataLoader(self.predict, batch_size=self.batch_size)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "JaDtzELEGPa-"
   },
   "source": [
    "---\n",
    "**Generate training data and train the model**\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "id": "pGa7w8bMwcbB"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Indexed 9 image sets...\n",
      "Processing set \"brow\"\n",
      "Processing set \"eyeBlinkLeft\"\n",
      "Processing set \"eyeBlinkRight\"\n",
      "Processing set \"take_fwd_expr\"\n",
      "Processing set \"take_fwd_text\"\n",
      "Processing set \"take_rnd_text\"\n",
      "Processing set \"take_eyes\"\n",
      "Processing set \"take_nose\"\n",
      "Processing set \"take_straight\"\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "      <iframe id=\"tensorboard-frame-95722d55ed2e40ac\" width=\"100%\" height=\"800\" frameborder=\"0\">\n",
       "      </iframe>\n",
       "      <script>\n",
       "        (function() {\n",
       "          const frame = document.getElementById(\"tensorboard-frame-95722d55ed2e40ac\");\n",
       "          const url = new URL(\"/\", window.location);\n",
       "          const port = 6006;\n",
       "          if (port) {\n",
       "            url.port = port;\n",
       "          }\n",
       "          frame.src = url;\n",
       "        })();\n",
       "      </script>\n",
       "    "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Applications\\Code\\miniconda3\\envs\\colab_env\\lib\\site-packages\\pytorch_lightning\\trainer\\connectors\\accelerator_connector.py:447: LightningDeprecationWarning: Setting `Trainer(gpus=1)` is deprecated in v1.7 and will be removed in v2.0. Please use `Trainer(accelerator='gpu', devices=1)` instead.\n",
      "  f\"Setting `Trainer(gpus={gpus!r})` is deprecated in v1.7 and will be removed\"\n",
      "GPU available: True (cuda), used: True\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "HPU available: False, using: 0 HPUs\n",
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n",
      "\n",
      "  | Name  | Type       | Params\n",
      "-------------------------------------\n",
      "0 | model | Sequential | 177 K \n",
      "-------------------------------------\n",
      "177 K     Trainable params\n",
      "0         Non-trainable params\n",
      "177 K     Total params\n",
      "0.708     Total estimated model params size (MB)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Sanity Checking: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Applications\\Code\\miniconda3\\envs\\colab_env\\lib\\site-packages\\pytorch_lightning\\trainer\\connectors\\data_connector.py:229: PossibleUserWarning: The dataloader, val_dataloader 0, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 16 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n",
      "  category=PossibleUserWarning,\n",
      "C:\\Applications\\Code\\miniconda3\\envs\\colab_env\\lib\\site-packages\\pytorch_lightning\\trainer\\connectors\\data_connector.py:229: PossibleUserWarning: The dataloader, train_dataloader, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 16 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n",
      "  category=PossibleUserWarning,\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "de5fe14e45c04de4a4cdff10805b93d5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "`Trainer.fit` stopped: `max_epochs=75` reached.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Label:\n",
      "[0.04 0.1  0.04 0.   0.   0.   0.   0.04 0.1  0.   0.   0.   0.   0.\n",
      " 0.   0.   0.   0.   0.   0.   0.   0.   0.01 0.01 0.03 0.   0.   0.\n",
      " 0.   0.   0.   0.   0.   0.   0.   0.   0.   0.   0.   0.   0.   0.02\n",
      " 0.01 0.   0.   0.   0.01 0.   0.   0.   0.   0.  ]\n",
      "\n",
      "Pred:\n",
      "[0.04 0.06 0.06 0.   0.   0.   0.   0.04 0.06 0.01 0.01 0.   0.   0.\n",
      " 0.   0.   0.   0.   0.   0.01 0.   0.   0.01 0.01 0.02 0.   0.   0.01\n",
      " 0.01 0.   0.   0.01 0.   0.03 0.01 0.01 0.01 0.03 0.02 0.02 0.01 0.\n",
      " 0.01 0.   0.   0.   0.01 0.   0.   0.01 0.02 0.  ]\n",
      "\n",
      "Abs error:\n",
      "[0.   0.05 0.02 0.   0.   0.   0.   0.   0.05 0.01 0.01 0.   0.   0.\n",
      " 0.   0.   0.   0.   0.   0.01 0.   0.   0.   0.   0.   0.   0.   0.\n",
      " 0.01 0.   0.   0.01 0.   0.03 0.01 0.01 0.01 0.03 0.02 0.02 0.01 0.01\n",
      " 0.   0.   0.   0.   0.01 0.   0.   0.01 0.02 0.  ]\n",
      "\n",
      "Max abs error: 0.04787\n",
      "\n",
      "MSE: 0.000175\n"
     ]
    }
   ],
   "source": [
    "# ---------------------------------------------------------------------------------------\n",
    "# Generate training set or load cached\n",
    "# ---------------------------------------------------------------------------------------\n",
    "if RUN_MEDIAPIPE_GEN:\n",
    "    # use mediapipe to generate input vectors\n",
    "    sets = []\n",
    "    with open(DATA_PATH + 'sets.txt', mode='r') as file:\n",
    "        sets = [{\n",
    "            'name': l.strip(),\n",
    "            'root': os.path.join(DATA_PATH, l.strip()),\n",
    "            'input': [],\n",
    "            'output': pd.read_csv(os.path.join(os.path.join(DATA_PATH, l.strip()), l.strip() + '.csv')),\n",
    "            'path': os.path.join(os.path.join(DATA_PATH, l.strip()), (l.strip() + '_{:04d}.jpg')),\n",
    "            'neutral_path': os.path.join(os.path.join(DATA_PATH, l.strip()), 'neutral.jpg'),\n",
    "            'count': len(glob.glob1(os.path.join(DATA_PATH, l.strip()),\"*.jpg\")) - 1, # subtracting one for neutral.jpg\n",
    "            } for l in file.readlines()]\n",
    "    print('Indexed ' + str(len(sets)) + ' image sets...')\n",
    "    # print(sets[0]['csv'].iloc[0])\n",
    "\n",
    "    with mp_face_mesh.FaceMesh(\n",
    "            static_image_mode=True,\n",
    "            refine_landmarks=REFINE_LANDMARKS,\n",
    "            max_num_faces=1,\n",
    "            min_detection_confidence=0.5) as face_mesh:\n",
    "\n",
    "        for set in sets:\n",
    "            print('Processing set \"' + set['name'] + '\"')\n",
    "\n",
    "            set['output'] = set['output'].to_numpy()\n",
    "\n",
    "            # Remove head orientation\n",
    "            while set['output'].shape[1] > 52:\n",
    "                set['output'] = np.delete(set['output'], 52, 1)\n",
    "\n",
    "            neutral = get_normalized_landmarks(set['neutral_path'])\n",
    "\n",
    "            for index in range(set['count']):\n",
    "                path = set['path'].format(index)\n",
    "                image = cv2.imread(path)\n",
    "                results = face_mesh.process(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))\n",
    "                input = convert_landmarks_to_model_input(results.multi_face_landmarks, neutral)\n",
    "                set['input'].append(input)\n",
    "\n",
    "        with open(os.path.join(DATA_PATH, 'sets.pickle'), 'wb') as f:\n",
    "            pickle.dump(sets, f)\n",
    "else:\n",
    "    with open(os.path.join(DATA_PATH, 'sets.pickle'), 'rb') as f:\n",
    "        sets = pickle.load(f)\n",
    "\n",
    "# ---------------------------------------------------------------------------------------\n",
    "# Train model or load cached\n",
    "# ---------------------------------------------------------------------------------------\n",
    "data_module = ARKitDataModule(sets)\n",
    "landmark_to_blendshape = FullyConnectedModel(NN_INPUT_SIZE, NN_OUTPUT_SIZE, [NN_HIDDEN_SIZE])\n",
    "\n",
    "model_pt_path = os.path.join(DATA_PATH, 'model.pt')\n",
    "model_onnx_path = os.path.join(DATA_PATH, 'model.onnx')\n",
    "\n",
    "dummy_input = torch.zeros(NN_INPUT_SIZE).float().reshape(NN_INPUT_SIZE, 1).transpose(0, 1)\n",
    "\n",
    "if RUN_TRAINING:\n",
    "    %load_ext tensorboard\n",
    "    %tensorboard --logdir lightning_logs/\n",
    "\n",
    "    # train\n",
    "    torch.manual_seed(1) \n",
    "    if torch.cuda.is_available(): # if you have GPUs\n",
    "        trainer = pl.Trainer(max_epochs=NN_EPOCHS, gpus=1)\n",
    "    else:\n",
    "        trainer = pl.Trainer(max_epochs=NN_EPOCHS)\n",
    "    trainer.fit(model=landmark_to_blendshape.float(), datamodule=data_module)\n",
    "    # optimize (TODO)\n",
    "    # save to pytorch and onnx formats\n",
    "    torch.save(landmark_to_blendshape.model.state_dict(), model_pt_path)\n",
    "    torch.onnx.export(landmark_to_blendshape.model, dummy_input, model_onnx_path, verbose=True)\n",
    "else:\n",
    "    # load pretrained\n",
    "    landmark_to_blendshape.model.load_state_dict(torch.load(model_pt_path))\n",
    "\n",
    "landmark_to_blendshape.model.eval()\n",
    "\n",
    "# ---------------------------------------------------------------------------------------\n",
    "# Test\n",
    "# ---------------------------------------------------------------------------------------\n",
    "def predict(name, index):\n",
    "    for s in sets:\n",
    "        if s['name'] != name:\n",
    "            continue;\n",
    "\n",
    "        x = s['input'][index]\n",
    "        l = s['output'][index]\n",
    "        x = torch.from_numpy(np.array(x).flatten()).float().reshape(NN_INPUT_SIZE, 1).transpose(0, 1)\n",
    "        y = np.clip(landmark_to_blendshape(x)[0].tolist(), 0, 1)\n",
    "        # print(\"Input:\")\n",
    "        # print(str(x))\n",
    "        print(\"\\nLabel:\")\n",
    "        print(str(np.around(l, 2)))\n",
    "        print(\"\\nPred:\")\n",
    "        print(str(np.around(y, 2)))\n",
    "        print(\"\\nAbs error:\")\n",
    "        print(str(np.around(np.abs(np.array(y) - np.array(l)), 2)))\n",
    "        print(\"\\nMax abs error: \" + str(np.around(np.abs(np.array(y) - np.array(l)).max(), 6)))\n",
    "        print(\"\\nMSE: \" + str(np.around(np.square(np.array(y) - np.array(l)).mean(), 6)))\n",
    "\n",
    "predict(sets[0]['name'], 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "provenance": []
  },
  "gpuClass": "premium",
  "kernelspec": {
   "display_name": "asr_env",
   "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.7.13"
  },
  "vscode": {
   "interpreter": {
    "hash": "d153c6b805ca38e125df211d9c3a3cc4d396b9f5bfb935b88898bae649c5b57e"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}