MoveNet Refined.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 0. Install and Import Dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install tensorflow==2.4.1 tensorflow-gpu==2.4.1 opencv-python matplotlib"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\shrie\\Anaconda3\\lib\\site-packages\\tensorboard\\compat\\tensorflow_stub\\dtypes.py:541: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_qint8 = np.dtype([(\"qint8\", np.int8, 1)])\n",
      "C:\\Users\\shrie\\Anaconda3\\lib\\site-packages\\tensorboard\\compat\\tensorflow_stub\\dtypes.py:542: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_quint8 = np.dtype([(\"quint8\", np.uint8, 1)])\n",
      "C:\\Users\\shrie\\Anaconda3\\lib\\site-packages\\tensorboard\\compat\\tensorflow_stub\\dtypes.py:543: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_qint16 = np.dtype([(\"qint16\", np.int16, 1)])\n",
      "C:\\Users\\shrie\\Anaconda3\\lib\\site-packages\\tensorboard\\compat\\tensorflow_stub\\dtypes.py:544: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_quint16 = np.dtype([(\"quint16\", np.uint16, 1)])\n",
      "C:\\Users\\shrie\\Anaconda3\\lib\\site-packages\\tensorboard\\compat\\tensorflow_stub\\dtypes.py:545: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_qint32 = np.dtype([(\"qint32\", np.int32, 1)])\n",
      "C:\\Users\\shrie\\Anaconda3\\lib\\site-packages\\tensorboard\\compat\\tensorflow_stub\\dtypes.py:550: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  np_resource = np.dtype([(\"resource\", np.ubyte, 1)])\n"
     ]
    }
   ],
   "source": [
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "from matplotlib import pyplot as plt\n",
    "import cv2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1. Load Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "interpreter = tf.lite.Interpreter(model_path='lite-model_movenet_singlepose_lightning_3.tflite')\n",
    "interpreter.allocate_tensors()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2. Make Detections"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n",
      "(1, 192, 192, 3)\n"
     ]
    }
   ],
   "source": [
    "cap = cv2.VideoCapture(0)\n",
    "while cap.isOpened():\n",
    "    ret, frame = cap.read()\n",
    "    #print(ret)\n",
    "    # Reshape image\n",
    "    img = frame.copy()\n",
    "    img = tf.image.resize_with_pad(np.expand_dims(img, axis=0), 192,192)\n",
    "    #print(img.shape)\n",
    "    input_image = tf.cast(img, dtype=tf.float32)\n",
    "    print(input_image.shape)\n",
    "    # Setup input and output \n",
    "    input_details = interpreter.get_input_details()\n",
    "    output_details = interpreter.get_output_details()\n",
    "    \n",
    "    # Make predictions \n",
    "    interpreter.set_tensor(input_details[0]['index'], np.array(input_image))\n",
    "    interpreter.invoke()\n",
    "    keypoints_with_scores = interpreter.get_tensor(output_details[0]['index'])\n",
    "    \n",
    "    # Rendering \n",
    "    draw_connections(frame, keypoints_with_scores, EDGES, 0.4)\n",
    "    draw_keypoints(frame, keypoints_with_scores, 0.4)\n",
    "    \n",
    "    cv2.imshow('MoveNet Lightning', frame)\n",
    "    \n",
    "    if cv2.waitKey(10) & 0xFF==ord('q'):\n",
    "        break\n",
    "        \n",
    "cap.release()\n",
    "cv2.destroyAllWindows()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3. Draw Keypoints"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def draw_keypoints(frame, keypoints, confidence_threshold):\n",
    "    y, x, c = frame.shape\n",
    "    shaped = np.squeeze(np.multiply(keypoints, [y,x,1]))\n",
    "    \n",
    "    for kp in shaped:\n",
    "        ky, kx, kp_conf = kp\n",
    "        if kp_conf > confidence_threshold:\n",
    "            cv2.circle(frame, (int(kx), int(ky)), 4, (0,255,0), -1) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 4. Draw Edges"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "EDGES = {\n",
    "    (0, 1): 'm',\n",
    "    (0, 2): 'c',\n",
    "    (1, 3): 'm',\n",
    "    (2, 4): 'c',\n",
    "    (0, 5): 'm',\n",
    "    (0, 6): 'c',\n",
    "    (5, 7): 'm',\n",
    "    (7, 9): 'm',\n",
    "    (6, 8): 'c',\n",
    "    (8, 10): 'c',\n",
    "    (5, 6): 'y',\n",
    "    (5, 11): 'm',\n",
    "    (6, 12): 'c',\n",
    "    (11, 12): 'y',\n",
    "    (11, 13): 'm',\n",
    "    (13, 15): 'm',\n",
    "    (12, 14): 'c',\n",
    "    (14, 16): 'c'\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def draw_connections(frame, keypoints, edges, confidence_threshold):\n",
    "    y, x, c = frame.shape\n",
    "    shaped = np.squeeze(np.multiply(keypoints, [y,x,1]))\n",
    "    \n",
    "    for edge, color in edges.items():\n",
    "        p1, p2 = edge\n",
    "        y1, x1, c1 = shaped[p1]\n",
    "        y2, x2, c2 = shaped[p2]\n",
    "        \n",
    "        if (c1 > confidence_threshold) & (c2 > confidence_threshold):      \n",
    "            cv2.line(frame, (int(x1), int(y1)), (int(x2), int(y2)), (0,0,255), 2)"
   ]
  },
  {
   "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}