demo-movenet-webcam.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e4ee3173-e406-4b78-b158-81b8a05c9c94",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import tensorflow_hub as hub\n",
    "import cv2\n",
    "from matplotlib import pyplot as plt\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a3b694f-5043-47bc-aa44-a196b399099e",
   "metadata": {},
   "source": [
    "Import Model from tf_hub"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "5efed31e-dad9-43c1-b013-615c0496f1b1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Metal device set to: Apple M1 Pro\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2022-06-24 08:43:35.086929: I tensorflow/core/common_runtime/pluggable_device/pluggable_device_factory.cc:305] Could not identify NUMA node of platform GPU ID 0, defaulting to 0. Your kernel may not have been built with NUMA support.\n",
      "2022-06-24 08:43:35.087087: I tensorflow/core/common_runtime/pluggable_device/pluggable_device_factory.cc:271] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 0 MB memory) -> physical PluggableDevice (device: 0, name: METAL, pci bus id: <undefined>)\n"
     ]
    }
   ],
   "source": [
    "model = hub.load(\"https://tfhub.dev/google/movenet/multipose/lightning/1\")\n",
    "movenet = model.signatures['serving_default']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "2675e124-c943-4e5a-acfb-a1df58e587ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to loop through each person detected and render\n",
    "def loop_through_people(frame, keypoints_with_scores, edges, confidence_threshold):\n",
    "    for person in keypoints_with_scores:\n",
    "        draw_connections(frame, person, edges, confidence_threshold)\n",
    "        draw_keypoints(frame, person, confidence_threshold)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "ef0dbfa9-871f-4780-88a1-f32216edd731",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Draw Landmarks\n",
    "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)), 6, (0,255,0), -1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "b2ce781c-2519-41e9-8216-9fa8f3547325",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Edges with corresponding Colors\n",
    "EDGES = {\n",
    "    (0, 1): (255,0,0),\n",
    "    (0, 2): (0,255,0),\n",
    "    (1, 3): (0,0,255),\n",
    "    (2, 4): (255,255,0),\n",
    "    (0, 5): (255,0,255),\n",
    "    (0, 6): (0,255,255),\n",
    "    (5, 7): (0,0,0),\n",
    "    (7, 9): (255,255,255),\n",
    "    (6, 8): (255,0,0),\n",
    "    (8, 10): (0,255,0),\n",
    "    (5, 6): (0,0,255),\n",
    "    (5, 11): (255,255,0),\n",
    "    (6, 12): (255,0,255),\n",
    "    (11, 12): (0,255,255),\n",
    "    (11, 13): (0,0,0),\n",
    "    (13, 15): (255,255,255),\n",
    "    (12, 14): (255,0,0),\n",
    "    (14, 16): (0,255,0)\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "7621c725-e708-40e6-a5f1-7076327c6456",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Draw Lines between Landmarks\n",
    "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)), color, 4)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20f5b349-9503-43bb-b435-af0749a65219",
   "metadata": {},
   "source": [
    "Webcam Capture (Hold 'q' (or press 3 times) to stop capture, close window and finish cell)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "818f4288-fccc-46dd-932b-7fc286fa1a1f",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2022-06-24 08:43:56.272013: W tensorflow/core/platform/profile_utils/cpu_utils.cc:128] Failed to get CPU frequency: 0 Hz\n",
      "2022-06-24 08:43:56.274136: I tensorflow/core/grappler/optimizers/custom_graph_optimizer_registry.cc:113] Plugin optimizer for device_type GPU is enabled.\n"
     ]
    }
   ],
   "source": [
    "cap = cv2.VideoCapture(0)\n",
    "#cap.set(3,1920) # adjust width\n",
    "#cap.set(4,1080) # adjust height\n",
    "\n",
    "\n",
    "while cap.isOpened():\n",
    "    ret, frame = cap.read()\n",
    "    \n",
    "    # Resize image\n",
    "    img = frame.copy()\n",
    "    img = tf.image.resize_with_pad(tf.expand_dims(img, axis=0), 384,640)\n",
    "    input_img = tf.cast(img, dtype=tf.int32)\n",
    "    \n",
    "    # Detection section\n",
    "    results = movenet(input_img)\n",
    "    keypoints_with_scores = results['output_0'].numpy()[:,:,:51].reshape((6,17,3))\n",
    "    \n",
    "    # Render keypoints \n",
    "    loop_through_people(frame, keypoints_with_scores, EDGES, 0.3)\n",
    "    \n",
    "    cv2.imshow('Movenet Multipose', frame)\n",
    "    \n",
    "    # Hold or Press 3 Times to finish\n",
    "    if cv2.waitKey(10) & 0xFF == ord('q'):\n",
    "        cap.release()\n",
    "        cv2.waitKey(0)\n",
    "        cv2.destroyAllWindows()\n",
    "        break\n",
    "        \n",
    "cap.release()\n",
    "cv2.waitKey(0)\n",
    "cv2.destroyAllWindows()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9 (ml-ds)",
   "language": "python",
   "name": "ml-ds"
  },
  "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.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}