06_GoogleMediaPipe.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Google MediaPipe\n",
    "Googles [MediaPipe](https://google.github.io/mediapipe/) is currently (as of 26.12.2020) in alpha state.\n",
    "\n",
    "It can be installed on Jetson Nano by following [given instructions](https://google.github.io/mediapipe/getting_started/install.html#installing-on-debian-and-ubuntu)\n",
    ". Python package needs to be build on Jetson Nano, following [these instructions](https://google.github.io/mediapipe/getting_started/python.html#building-mediapipe-python-package)\n",
    "\n",
    "TLDR:\n",
    "```bash\n",
    "# First install Bazel 3.4.1\n",
    "mkdir $HOME/bazel-3.4.1\n",
    "cd $HOME/bazel-3.4.1\n",
    "wget https://github.com/bazelbuild/bazel/releases/download/3.4.1/bazel-3.4.1-dist.zip\n",
    "sudo apt-get install build-essential openjdk-8-jdk python zip unzip\n",
    "unzip bazel-3.4.1-dist.zip\n",
    "env EXTRA_BAZEL_ARGS=\"--host_javabase=@local_jdk//:jdk\" bash ./compile.sh\n",
    "sudo cp output/bazel /usr/local/bin/\n",
    "\n",
    "# Then the MediaPipe\n",
    "cd $HOME\n",
    "git clone https://github.com/google/mediapipe.git\n",
    "cd mediapipe\n",
    "# if you have not updated OpenCV to 4.X (compiled from source)\n",
    "sed -i \"s/x86_64-linux-gnu/aarch64-linux-gnu/g\" third_party/opencv_linux.BUILD\n",
    "# else\n",
    "sed -i \"s/x86_64-linux-gnu//g\" third_party/opencv_linux.BUILD\n",
    "sed -i \"s/include/include\\/opencv4/g\" third_part/opencv_linux.BUILD\n",
    "# then\n",
    "sudo -H pip3 install -r requirements.txt\n",
    "python3 setup.py gen_protos\n",
    "sudo -H python3 setup.py install --link-opencv\n",
    "```\n",
    "\n",
    "Performance of this model is moderate slow\n",
    "```\n",
    "holistic ~ 3.5 fps\n",
    "hands ~ 8.5 fps\n",
    "pose ~ 6.5 fps\n",
    "face ~ 14 fps\n",
    "```\n",
    "\n",
    "Memory wise this model is superb to others I have tested sofar whole model taking only ~ 1GB in total\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "import traitlets\n",
    "\n",
    "import numpy as np\n",
    "import mediapipe as mp\n",
    "\n",
    "from jetutils import SimpleTimer\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mp_drawing = mp.solutions.drawing_utils\n",
    "class MediaPipeAdapter(traitlets.HasTraits):\n",
    "    output_frame = traitlets.Any()\n",
    "    draw_original = traitlets.Bool(default_value=True)\n",
    "    solution_name = traitlets.Unicode(default_value='holistic')\n",
    "    \n",
    "    def __init__(self):\n",
    "        self._timer = SimpleTimer()\n",
    "        self._full_timer = SimpleTimer()\n",
    "        self._model = {\n",
    "            'holistic': mp.solutions.holistic.Holistic(smooth_landmarks=True,\n",
    "                                                       min_detection_confidence=0.4,\n",
    "                                                       min_tracking_confidence=0.5),\n",
    "            'hands': mp.solutions.hands.Hands(max_num_hands=2,\n",
    "                                              min_detection_confidence=0.4,\n",
    "                                              min_tracking_confidence=0.5),\n",
    "            'face': mp.solutions.face_mesh.FaceMesh(min_detection_confidence=0.4,\n",
    "                                                    min_tracking_confidence=0.5),\n",
    "            'pose': mp.solutions.pose.Pose(min_detection_confidence=0.4,\n",
    "                                           min_tracking_confidence=0.5),\n",
    "        }\n",
    "        self._draw = {\n",
    "            'holistic': self._draw_holistic,\n",
    "            'hands': self._draw_hands,\n",
    "            'face': self._draw_face,\n",
    "            'pose': self._draw_pose,\n",
    "        }\n",
    "\n",
    "        self._drawing_spec = mp_drawing.DrawingSpec(thickness=1, circle_radius=1)\n",
    "        self.output_frame = np.zeros((480, 640, 3), dtype=np.uint8)\n",
    "        \n",
    "    def __del__(self):\n",
    "        for model in self._model:\n",
    "            self._model[model].close()\n",
    "        \n",
    "    def _draw_holistic(self, frame, results):\n",
    "        mp_drawing.draw_landmarks(\n",
    "            frame, results.face_landmarks, mp.solutions.holistic.FACE_CONNECTIONS)\n",
    "        mp_drawing.draw_landmarks(\n",
    "            frame, results.left_hand_landmarks, mp.solutions.holistic.HAND_CONNECTIONS)\n",
    "        mp_drawing.draw_landmarks(\n",
    "            frame, results.right_hand_landmarks, mp.solutions.holistic.HAND_CONNECTIONS)\n",
    "        mp_drawing.draw_landmarks(\n",
    "            frame, results.pose_landmarks, mp.solutions.holistic.POSE_CONNECTIONS)\n",
    "    \n",
    "    def _draw_hands(self, frame, results):\n",
    "        if results.multi_hand_landmarks:\n",
    "            for hand_landmarks in results.multi_hand_landmarks:\n",
    "                mp_drawing.draw_landmarks(\n",
    "                    frame, hand_landmarks, mp.solutions.hands.HAND_CONNECTIONS)\n",
    "\n",
    "    def _draw_face(self, frame, results):\n",
    "        if results.multi_face_landmarks:\n",
    "            for face_landmarks in results.multi_face_landmarks:\n",
    "                mp_drawing.draw_landmarks(frame, face_landmarks,\n",
    "                                          mp.solutions.face_mesh.FACE_CONNECTIONS,\n",
    "                                          landmark_drawing_spec=self._drawing_spec,\n",
    "                                          connection_drawing_spec=self._drawing_spec)\n",
    "\n",
    "    def _draw_pose(self, frame, results):\n",
    "        mp_drawing.draw_landmarks(frame,\n",
    "                                  results.pose_landmarks, mp.solutions.pose.POSE_CONNECTIONS)\n",
    "    \n",
    "    def _dbgstr(self, results):\n",
    "        # done very ugly here...\n",
    "        if self.solution_name == 'holistic':\n",
    "            dbgstr = 'Landmarks:\\n'\n",
    "            dbgstr += 'face: {} landmarks\\n\\n'.format(len(results.face_landmarks.landmark)) if results.face_landmarks else 'face: NONE\\n\\n'\n",
    "            if results.pose_landmarks:\n",
    "                for i in range(len(results.pose_landmarks.landmark)):\n",
    "                    dbgstr += '{}{} x:{}\\t y:{}\\t visibility %:{}\\n'.format(mp.solutions.holistic.PoseLandmark(i).name.lower(),\n",
    "                                                                            ' '*(20-len(mp.solutions.holistic.PoseLandmark(i).name)),\n",
    "                                                                            int(results.pose_landmarks.landmark[i].x*640),\n",
    "                                                                            int(results.pose_landmarks.landmark[i].y*480),\n",
    "                                                                            int(results.pose_landmarks.landmark[i].visibility*100))\n",
    "                dbgstr += '\\n\\n'\n",
    "            if results.left_hand_landmarks:\n",
    "                dbgstr += 'left hand:\\n'\n",
    "                for i in range(len(results.left_hand_landmarks.landmark)):\n",
    "                    dbgstr += '{}{} x:{}\\t y:{}\\t z(depth):{}\\n'.format(mp.solutions.holistic.HandLandmark(i).name.lower(),\n",
    "                                                                 ' '*(20-len(mp.solutions.holistic.HandLandmark(i).name)),\n",
    "                                                                 int(results.left_hand_landmarks.landmark[i].x*640),\n",
    "                                                                 int(results.left_hand_landmarks.landmark[i].y*480),\n",
    "                                                                 int(results.left_hand_landmarks.landmark[i].z*640))\n",
    "                dbgstr += '\\n\\n'\n",
    "            else:\n",
    "                dbgstr += 'left hand: NONE\\n\\n'\n",
    "\n",
    "            if results.right_hand_landmarks:\n",
    "                dbgstr += 'right hand:\\n'\n",
    "                for i in range(len(results.right_hand_landmarks.landmark)):\n",
    "                    dbgstr += '{}{} x:{}\\t y:{}\\t z(depth):{}\\n'.format(mp.solutions.holistic.HandLandmark(i).name.lower(),\n",
    "                                                                 ' '*(20-len(mp.solutions.holistic.HandLandmark(i).name)),\n",
    "                                                                 int(results.right_hand_landmarks.landmark[i].x*640),\n",
    "                                                                 int(results.right_hand_landmarks.landmark[i].y*480),\n",
    "                                                                 int(results.right_hand_landmarks.landmark[i].z*640))\n",
    "                dbgstr += '\\n\\n'\n",
    "            else:\n",
    "                dbgstr += 'right hand: NONE\\n\\n'\n",
    "\n",
    "        elif self.solution_name == 'hands':\n",
    "            dbgstr = 'hands {}\\n\\n'.format(len(results.multi_hand_landmarks)) if results.multi_hand_landmarks else 'hands NONE\\n'\n",
    "            if results.multi_hand_landmarks:\n",
    "                hnd = 0\n",
    "                for hand_landmarks in results.multi_hand_landmarks:\n",
    "                    dbgstr += 'hand {}\\n'.format(hnd)\n",
    "                    for i in range(len(hand_landmarks.landmark)):\n",
    "                        dbgstr += '{}{} x:{}\\t y:{}\\t z(depth):{}\\n'.format(mp.solutions.holistic.HandLandmark(i).name.lower(),\n",
    "                                                                     ' '*(20-len(mp.solutions.holistic.HandLandmark(i).name)),\n",
    "                                                                     int(hand_landmarks.landmark[i].x*640),\n",
    "                                                                     int(hand_landmarks.landmark[i].y*480),\n",
    "                                                                     int(hand_landmarks.landmark[i].z*640))\n",
    "                    dbgstr += '\\n\\n'\n",
    "                    hnd += 1\n",
    "            \n",
    "        elif self.solution_name == 'face':\n",
    "            dbgstr = 'faces {}\\n'.format(len(results.multi_face_landmarks)) if results.multi_face_landmarks else 'faces NONE\\n'\n",
    "                \n",
    "        elif self.solution_name == 'pose':\n",
    "            dbgstr = 'Landmarks:\\n'\n",
    "            if results.pose_landmarks:\n",
    "                for i in range(len(results.pose_landmarks.landmark)):\n",
    "                    dbgstr += '{}{} x:{}\\t y:{}\\t visibility %:{}\\n'.format(mp.solutions.pose.PoseLandmark(i).name.lower(),\n",
    "                                                                            ' '*(20-len(mp.solutions.pose.PoseLandmark(i).name)),\n",
    "                                                                            int(results.pose_landmarks.landmark[i].x*640),\n",
    "                                                                            int(results.pose_landmarks.landmark[i].y*480),\n",
    "                                                                            int(results.pose_landmarks.landmark[i].visibility*100))\n",
    "        else:\n",
    "            dbgstr = ''\n",
    "        return dbgstr\n",
    "\n",
    "            \n",
    "    def process_frame(self, frame):\n",
    "        frame = frame.astype(np.uint8)\n",
    "        if self.draw_original:\n",
    "            outframe = frame.copy()\n",
    "        else:\n",
    "            outframe = np.zeros_like(frame)\n",
    "                \n",
    "        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)\n",
    "        frame.flags.writeable = False #pass reference - should be faster\n",
    "\n",
    "        with self._full_timer: \n",
    "            with self._timer:\n",
    "                results = self._model[str(self.solution_name)].process(frame)\n",
    "            model_time = self._timer.time\n",
    "            frame.flags.writeable = True\n",
    "            with self._timer:\n",
    "                self._draw[str(self.solution_name)](outframe, results)\n",
    "            draw_time = self._timer.time\n",
    "        self.output_frame = outframe\n",
    "        return 'model {}\\nfps {}\\nmodel time {}\\ndraw time {}\\n\\n{}'.format(self.solution_name,\n",
    "                                                                            self._full_timer.fps,\n",
    "                                                                            model_time,\n",
    "                                                                            draw_time,\n",
    "                                                                            self._dbgstr(results))\n",
    "\n",
    "mediapipe = MediaPipeAdapter()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from jetutils import GstCamera, bgr8_to_jpeg\n",
    "camera = GstCamera()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test all models once\n",
    "for model in mediapipe._model:\n",
    "    mediapipe.solution_name = value = model\n",
    "    print(mediapipe.process_frame(camera.read()))\n",
    "    print('-'*50)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ipywidgets\n",
    "from IPython.display import display\n",
    "from sidecar import Sidecar\n",
    "\n",
    "image_original = ipywidgets.Image(format='jpeg', width=camera.width, height=camera.height)\n",
    "image_processed = ipywidgets.Image(format='jpeg', width=camera.width, height=camera.height)\n",
    "image_original.value = bgr8_to_jpeg(np.zeros((camera.height, camera.width, 3), dtype=np.uint8))\n",
    "image_processed.value = bgr8_to_jpeg(np.zeros((camera.height, camera.width, 3), dtype=np.uint8))\n",
    "\n",
    "\n",
    "debug_out = ipywidgets.Textarea(value='',\n",
    "                                disabled=True,\n",
    "                                layout=ipywidgets.Layout(width='640px', height='1200px'))\n",
    "images_out = ipywidgets.HBox([image_original, image_processed])\n",
    "select_outmode = ipywidgets.ToggleButtons(options=['holistic', 'hands', 'pose', 'face'],\n",
    "                                     value='holistic', description='model ',\n",
    "                                     style={'description_width': 'initial'})\n",
    "select_original = ipywidgets.ToggleButton(value=True, description='over orginal',\n",
    "                                         style={'description_wifth': 'initial'})\n",
    "\n",
    "control_box = ipywidgets.HBox([select_outmode, select_original])\n",
    "all_box = ipywidgets.VBox([images_out, control_box, debug_out])\n",
    "\n",
    "_sidecar = Sidecar(title='output')\n",
    "with _sidecar:\n",
    "    display(all_box)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "traitlets.dlink((camera, 'value'), (image_original, 'value'), transform=bgr8_to_jpeg)\n",
    "traitlets.dlink((mediapipe,'output_frame'), (image_processed, 'value'), transform=bgr8_to_jpeg)\n",
    "# traitlets.dlink((select_outmode, 'value'), (mediapipe, 'solution_name'))\n",
    "traitlets.dlink((select_original, 'value'), (mediapipe, 'draw_original'))\n",
    "\n",
    "def process(change):\n",
    "    debug_out.value = mediapipe.process_frame(change['new'])\n",
    "    \n",
    "def change_model(change):\n",
    "    # stop camera while changin mode to prevent occasionally hangs\n",
    "    camera.running = False\n",
    "    mediapipe.solution_name = change['new']\n",
    "    camera.running = True\n",
    "    \n",
    "camera.observe(process, names='value')\n",
    "select_outmode.observe(change_model, names='value')\n",
    "camera.running = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "camera.unobserve_all()\n",
    "camera.running = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "del mediapipe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "del camera"
   ]
  }
 ],
 "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}