{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "674b0082-6453-4804-84ac-3f48ca3a166e",
   "metadata": {},
   "source": [
    "# Initialize"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e50f3a42-c144-4501-b4bd-0ad7ff090458",
   "metadata": {},
   "source": [
    "## Import libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84c80a5a-22ce-449d-a72c-587845ec2a70",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import re\n",
    "import json\n",
    "import math\n",
    "import time\n",
    "import copy\n",
    "import io\n",
    "import random\n",
    "import datetime\n",
    "import threading\n",
    "import signal\n",
    "import traceback\n",
    "import cProfile"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "342ca5c2-3743-4c76-bad6-22364460f998",
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "import numpy as np\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9fee88c8-c314-4f7d-9706-e7bae5bb7009",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d4a3fde-2127-4f30-8503-d6d341f6dd6f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "import panoramasdk"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1734cd6-311c-4baa-8463-39e9c460f5a7",
   "metadata": {},
   "source": [
    "## Configure GPU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d05dde4-010f-46c1-a88e-98f8266af439",
   "metadata": {},
   "outputs": [],
   "source": [
    "gpus = tf.config.list_physical_devices('GPU')\n",
    "gpus"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "889837db-7ed1-4b7b-a260-fb0484997b1a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Don't allocate huge memory unnecessarily\n",
    "tf.config.experimental.set_memory_growth( gpus[0], True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2bfc9f0a-2e20-406d-8cae-450d35108b49",
   "metadata": {},
   "source": [
    "## Instantiate panoramasdk.node object"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a4747e0-4a41-4e01-a8dc-f81601c33d46",
   "metadata": {},
   "outputs": [],
   "source": [
    "node = panoramasdk.node()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f94dd0f4-2dbc-4369-891b-a3128beb4f10",
   "metadata": {},
   "source": [
    "# Hello World"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e30aef2-aecc-493c-800b-1620696c0928",
   "metadata": {},
   "source": [
    "## Get frames from camera and Put to HDMI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64e25eec-07a2-423e-b249-8ced294b2f64",
   "metadata": {},
   "outputs": [],
   "source": [
    "media_list = node.inputs.video_in.get()\n",
    "\n",
    "media_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d201de3-fdf6-40bd-8583-ef88fe1bb5b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "media_list[0].image.shape, media_list[0].image.dtype"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4244e7b-4ee0-4567-80a7-7b34f0e54bbd",
   "metadata": {},
   "outputs": [],
   "source": [
    "def previewImage( image ):\n",
    "    \n",
    "    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
    "    \n",
    "    plt.figure( figsize = ( 10, 10 ) )\n",
    "    plt.imshow( image_rgb, interpolation='antialiased' )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1506e665-64ed-4ec0-8644-b27d5cb30ff3",
   "metadata": {},
   "outputs": [],
   "source": [
    "previewImage(media_list[0].image)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c54bf4dd-3e9a-401b-b07d-08cd3afde5a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "node.outputs.video_out.put( media_list )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "076185bb-36cd-427a-add3-e5ec8fd08282",
   "metadata": {},
   "source": [
    "## Render \"Hello World\" text by OpenCV"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bdb03f1e-fade-4c10-b752-b52840eacf3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "text_color = (255,255,255)\n",
    "text_shadow_color = (0,0,0)\n",
    "text_thickness = 3\n",
    "text_shadow_thickness = 5\n",
    "text_scale = 4\n",
    "\n",
    "def renderHelloWorld( media_list ):\n",
    "\n",
    "    image = media_list[0].image\n",
    "    \n",
    "    h, w, _ = image.shape\n",
    "    \n",
    "    cv2.putText( image, f\"Hello World\", (22, h//2+2), fontFace=cv2.FONT_HERSHEY_PLAIN, fontScale=text_scale, color=text_shadow_color, thickness=text_shadow_thickness, lineType=cv2.LINE_AA )\n",
    "    cv2.putText( image, f\"Hello World\", (20, h//2), fontFace=cv2.FONT_HERSHEY_PLAIN, fontScale=text_scale, color=text_color, thickness=text_thickness, lineType=cv2.LINE_AA )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10bf94a2-1e6c-4f83-892e-ab30900b4f9e",
   "metadata": {},
   "outputs": [],
   "source": [
    "renderHelloWorld( media_list )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c2c6de2-612c-4dbd-b054-22d91d1b26ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "previewImage(media_list[0].image)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c3487e6-9854-4d70-8605-1751b3560864",
   "metadata": {},
   "source": [
    "## Main loop (for Hello World)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a37684e0-7087-4847-bea5-d0f4897b4948",
   "metadata": {},
   "outputs": [],
   "source": [
    "def mainLoop():\n",
    "    try:\n",
    "        while True:\n",
    "            \n",
    "            # get video frame(s) from camera(s)\n",
    "            media_list = node.inputs.video_in.get()\n",
    "            \n",
    "            # render Hello World Text in the image(s)\n",
    "            renderHelloWorld(media_list)\n",
    "            \n",
    "            # output to HDMI\n",
    "            node.outputs.video_out.put( media_list )\n",
    "            \n",
    "    except KeyboardInterrupt:\n",
    "        pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a69e5ae3-ab20-4f4a-9967-a1508d2ab695",
   "metadata": {},
   "outputs": [],
   "source": [
    "mainLoop()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cfb7ea75-e151-4fcc-a875-a6c44ad88c0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "!free"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18a8f7d1-50e2-4874-b28d-d7572edd16e4",
   "metadata": {},
   "source": [
    "# Pose estimation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "283eb994-47ba-4ddd-8ec0-6dd24f45f075",
   "metadata": {},
   "source": [
    "## Load model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "837fd7d9-a863-49c2-932b-e9ade08b220e",
   "metadata": {},
   "source": [
    "**Manual steps:**\n",
    "\n",
    "1. Prepare a model file.\n",
    "    * Option 1 : Download \"movenet_multipose_lightning_1.tar.gz\" from https://tfhub.dev/google/movenet/multipose/lightning/1 on your PC.\n",
    "    * Option 2 : Using TensorRT, quantize the movenet_multipose_lightning_1 model, and prepare movenet_multipose_lightning_1_trt_fp16.tar.gz.\n",
    "2. Upload the prepared model file onto Jupyter (You can Drag & drop the file to browser pane ).\n",
    "3. Open a Terminal on Jupyter and run following commands.\n",
    "\n",
    "    ``` sh\n",
    "    tar xvzf ../movenet_multipose_lightning_1.tar.gz\n",
    "    ```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b6e4e3a-31be-4f50-9f93-267aa6b3f9b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "movenet_model = tf.saved_model.load(\"./movenet_multipose_lightning_1\")\n",
    "#movenet_model = tf.saved_model.load(\"./movenet_multipose_lightning_1_trt_fp16\")\n",
    "\n",
    "movenet_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dcdd2745-1ce5-4c74-bd25-8d61b07a1827",
   "metadata": {},
   "outputs": [],
   "source": [
    "pose_estimator = movenet_model.signatures[\"serving_default\"]\n",
    "pose_estimator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63e9799d-55d3-41e2-b99c-15512864b074",
   "metadata": {},
   "outputs": [],
   "source": [
    "!free"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ade80dc1-3c15-4867-9839-be9dea8e582b",
   "metadata": {},
   "source": [
    "## Preprocess image and run pose estimation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a83c819-5067-4f70-8d3a-933b50cf2fb0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def estimatePose( image ):\n",
    "\n",
    "    input_resolution = ( 160, 256 )\n",
    "    \n",
    "    image = tf.expand_dims( image, axis=0 )\n",
    "\n",
    "    image = tf.image.resize( image, input_resolution )\n",
    "\n",
    "    # BGR to RGB\n",
    "    image = tf.reverse(image, axis=[-1])\n",
    "\n",
    "    image = tf.cast( image, dtype=tf.int32 )\n",
    "\n",
    "    result = pose_estimator(image)\n",
    "\n",
    "    boxes = result[\"output_0\"].numpy()[:,:,3*17:].reshape(6,5)\n",
    "    joints = result[\"output_0\"].numpy()[:,:,:3*17].reshape(6,17,3)\n",
    "    \n",
    "    return boxes, joints\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82f634e2-b074-4cf1-bbe9-1fdb83dad1f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "boxes, joints = estimatePose( media_list[0].image )\n",
    "\n",
    "boxes, joints"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "311e770b-5ac2-45ae-824e-7d685a419d3e",
   "metadata": {},
   "outputs": [],
   "source": [
    "!free"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d38b195c-f67c-4be4-954b-08569e0b62d8",
   "metadata": {},
   "source": [
    "## Render inference result with OpenCV"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b95fda2b-c43f-4983-852b-e3baf49f2641",
   "metadata": {},
   "outputs": [],
   "source": [
    "box_threshold = 0.1\n",
    "joint_threshold = 0.5\n",
    "box_color = (255,0,0)\n",
    "box_thickness = 2\n",
    "dot_color = (0,255,0)\n",
    "dot_size = 3\n",
    "\n",
    "def renderResult( image, boxes, joints_list ):\n",
    "    \n",
    "    h, w, _ = image.shape\n",
    "    \n",
    "    for box, joints in zip( boxes, joints_list ):\n",
    "\n",
    "        if box[4] < box_threshold:\n",
    "            continue\n",
    "        \n",
    "        box_in_camera_space = (\n",
    "            int( box[1].item() * w ),\n",
    "            int( box[0].item() * h ),\n",
    "            int( box[3].item() * w ),\n",
    "            int( box[2].item() * h ), \n",
    "        )\n",
    "\n",
    "        cv2.rectangle( \n",
    "            image, \n",
    "            box_in_camera_space[0:2], \n",
    "            box_in_camera_space[2:4], \n",
    "            color = box_color, thickness = box_thickness, lineType=cv2.LINE_8\n",
    "        )\n",
    "        \n",
    "        for joint in joints:\n",
    "            if joint[2] < joint_threshold:\n",
    "                continue\n",
    "            \n",
    "            joint_in_camera_space = (\n",
    "                int( joint[1].item() * w ) - dot_size,\n",
    "                int( joint[0].item() * h ) - dot_size,\n",
    "                int( joint[1].item() * w ) + dot_size,\n",
    "                int( joint[0].item() * h ) + dot_size, \n",
    "            )\n",
    "\n",
    "            cv2.rectangle( \n",
    "                image, \n",
    "                joint_in_camera_space[0:2], \n",
    "                joint_in_camera_space[2:4], \n",
    "                color = dot_color, thickness = -1\n",
    "            )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "053c4a08-6b82-4da1-a1b9-96184c0f066a",
   "metadata": {},
   "outputs": [],
   "source": [
    "renderResult( media_list[0].image, boxes, joints )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c7a156c-54a8-4c53-8d47-bd09cce3cafd",
   "metadata": {},
   "outputs": [],
   "source": [
    "previewImage(media_list[0].image)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fdbfadde-ad38-40b6-8170-f8de55bbe15b",
   "metadata": {},
   "source": [
    "## Main loop (for pose estimation)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3222bfad-0659-4182-8a57-49d89022dbd7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def mainLoop():\n",
    "    try:\n",
    "        while True:\n",
    "            media_list = node.inputs.video_in.get()\n",
    "            \n",
    "            for media_obj in media_list:\n",
    "                boxes, joints = estimatePose( media_obj.image )\n",
    "                renderResult( media_obj.image, boxes, joints )\n",
    "            \n",
    "            node.outputs.video_out.put( media_list )\n",
    "            \n",
    "    except KeyboardInterrupt:\n",
    "        pass\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64563e76-f13e-4d54-914e-cd2f7f716800",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "mainLoop()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "696b98d2-b51d-4aa8-a6b1-98f3498e4a36",
   "metadata": {},
   "outputs": [],
   "source": [
    "!free"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a8b3da4d-ab90-4cd8-9530-dcde6936bd10",
   "metadata": {},
   "source": [
    "# Tips"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e326c1ef-18d3-4145-bc5f-cf7efa7e81be",
   "metadata": {},
   "source": [
    "## Run the main loop with profiler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b08ef0b7-3f4c-4a6a-b0de-e82aa95e07cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "cProfile.runctx( \"mainLoop()\", globals(), locals() )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9a90a53-0910-4f3b-a25e-600a25f2ead0",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.5"
  },
  "toc-autonumbering": true,
  "toc-showmarkdowntxt": false,
  "toc-showtags": false
 },
 "nbformat": 4,
 "nbformat_minor": 5
}