trt_pose/live_demo_karate.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#Referenced from https://github.com/NVIDIA-AI-IOT/trt_pose"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, let's load the JSON file which describes the human pose task.  This is in COCO format, it is the category descriptor pulled from the annotations file.  We modify the COCO category slightly, to add a neck keypoint.  We will use this task description JSON to create a topology tensor, which is an intermediate data structure that describes the part linkages, as well as which channels in the part affinity field each linkage corresponds to."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import trt_pose.coco\n",
    "\n",
    "with open('human_pose.json', 'r') as f:\n",
    "    human_pose = json.load(f)\n",
    "\n",
    "topology = trt_pose.coco.coco_category_to_topology(human_pose)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we'll load our model.  Each model takes at least two parameters, *cmap_channels* and *paf_channels* corresponding to the number of heatmap channels\n",
    "and part affinity field channels.  The number of part affinity field channels is 2x the number of links, because each link has a channel corresponding to the\n",
    "x and y direction of the vector field for each link."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import trt_pose.models\n",
    "\n",
    "num_parts = len(human_pose['keypoints'])\n",
    "num_links = len(human_pose['skeleton'])\n",
    "\n",
    "model = trt_pose.models.resnet18_baseline_att(num_parts, 2 * num_links).cuda().eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, let's load the model weights.  You will need to download these according to the table in the README."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "MODEL_WEIGHTS = 'resnet18_baseline_att_224x224_A_epoch_249.pth'\n",
    "\n",
    "model.load_state_dict(torch.load(MODEL_WEIGHTS))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order to optimize with TensorRT using the python library *torch2trt* we'll also need to create some example data.  The dimensions\n",
    "of this data should match the dimensions that the network was trained with.  Since we're using the resnet18 variant that was trained on\n",
    "an input resolution of 224x224, we set the width and height to these dimensions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "WIDTH = 224\n",
    "HEIGHT = 224\n",
    "\n",
    "data = torch.zeros((1, 3, HEIGHT, WIDTH)).cuda()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we'll use [torch2trt](https://github.com/NVIDIA-AI-IOT/torch2trt) to optimize the model.  We'll enable fp16_mode to allow optimizations to use reduced half precision."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The optimized model may be saved so that we do not need to perform optimization again, we can just load the model.  Please note that TensorRT has device specific optimizations, so you can only use an optimized model on similar platforms."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch2trt\n",
    "\n",
    "model_trt = torch2trt.torch2trt(model, [data], fp16_mode=True, max_workspace_size=1<<25)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "OPTIMIZED_MODEL = 'resnet18_baseline_att_224x224_A_epoch_249_trt.pth'\n",
    "\n",
    "torch.save(model_trt.state_dict(), OPTIMIZED_MODEL)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We could then load the saved model using *torch2trt* as follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch2trt import TRTModule\n",
    "\n",
    "model_trt = TRTModule()\n",
    "model_trt.load_state_dict(torch.load(OPTIMIZED_MODEL))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can benchmark the model in FPS with the following code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "t0 = time.time()\n",
    "torch.cuda.current_stream().synchronize()\n",
    "for i in range(50):\n",
    "    y = model_trt(data)\n",
    "torch.cuda.current_stream().synchronize()\n",
    "t1 = time.time()\n",
    "\n",
    "print(50.0 / (t1 - t0))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, let's define a function that will preprocess the image, which is originally in BGR8 / HWC format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "import torchvision.transforms as transforms\n",
    "import PIL.Image\n",
    "\n",
    "mean = torch.Tensor([0.485, 0.456, 0.406]).cuda()\n",
    "std = torch.Tensor([0.229, 0.224, 0.225]).cuda()\n",
    "device = torch.device('cuda')\n",
    "\n",
    "def preprocess(image):\n",
    "    global device\n",
    "    device = torch.device('cuda')\n",
    "    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
    "    image = PIL.Image.fromarray(image)\n",
    "    image = transforms.functional.to_tensor(image).to(device)\n",
    "    image.sub_(mean[:, None, None]).div_(std[:, None, None])\n",
    "    return image[None, ...]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we'll define two callable classes that will be used to parse the objects from the neural network, as well as draw the parsed objects on an image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from trt_pose.draw_objects import DrawObjects\n",
    "from trt_pose.parse_objects import ParseObjects\n",
    "\n",
    "parse_objects = ParseObjects(topology)\n",
    "draw_objects = DrawObjects(topology)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Assuming you're using NVIDIA Jetson, you can use the [jetcam](https://github.com/NVIDIA-AI-IOT/jetcam) package to create an easy to use camera that will produce images in BGR8/HWC format.\n",
    "\n",
    "If you're not on Jetson, you may need to adapt the code below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from jetcam.usb_camera import USBCamera\n",
    "# from jetcam.csi_camera import CSICamera\n",
    "from jetcam.utils import bgr8_to_jpeg\n",
    "\n",
    "camera = USBCamera(width=WIDTH, height=HEIGHT, capture_fps=30)\n",
    "# camera = CSICamera(width=WIDTH, height=HEIGHT, capture_fps=30)\n",
    "\n",
    "camera.running = True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Connect to the joycontrol tcp server.\n",
    "\n",
    "Skip the following cell if the joycontrol tcp server is not running yet."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import asyncio\n",
    "loop = asyncio.get_event_loop()\n",
    "tcpReader, tcpWriter = await asyncio.open_connection('127.0.0.1', 8080, loop=loop)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Handler for sending joycontrol command manually."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def on_button_clicked(b):\n",
    "    global tcpWriter\n",
    "    tcpWriter.write(b.tooltip.encode())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "ipywidgets for user interactions.\n",
    "\n",
    "Check the \"Control with data\" checkbox for sending the game controls from processTrtData to the joycontrol tcp server."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we'll create a widget which will be used to display the camera feed with visualizations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "import ipywidgets\n",
    "from IPython.display import display\n",
    "\n",
    "image_w = ipywidgets.Image(format='jpeg')\n",
    "display(image_w)\n",
    "\n",
    "controllerCheckBox_w = ipywidgets.Checkbox(\n",
    "    value=False,\n",
    "    description='Control with data',\n",
    "    disabled=False,\n",
    "    indent=False\n",
    ")\n",
    "display(ipywidgets.HBox([controllerCheckBox_w]))\n",
    "\n",
    "dataText_w = ipywidgets.Text(\"Value\", description=\"Data\", layout=ipywidgets.Layout(width='100%', height='30px'))\n",
    "display(dataText_w)\n",
    "\n",
    "upButton = ipywidgets.Button(tooltip='u',icon='arrow-up')\n",
    "upButton.on_click(on_button_clicked)\n",
    "downButton = ipywidgets.Button(tooltip='d',icon='arrow-down')\n",
    "downButton.on_click(on_button_clicked)\n",
    "leftButton = ipywidgets.Button(tooltip='l',icon='arrow-left')\n",
    "leftButton.on_click(on_button_clicked)\n",
    "rightButton = ipywidgets.Button(tooltip='r',icon='arrow-right')\n",
    "rightButton.on_click(on_button_clicked)\n",
    "\n",
    "aButton = ipywidgets.Button(description=\"A\", tooltip='a')\n",
    "aButton.on_click(on_button_clicked)\n",
    "bButton = ipywidgets.Button(description=\"B\", tooltip='b')\n",
    "bButton.on_click(on_button_clicked)\n",
    "xButton = ipywidgets.Button(description=\"X\", tooltip='x')\n",
    "xButton.on_click(on_button_clicked)\n",
    "yButton = ipywidgets.Button(description=\"Y\", tooltip='y')\n",
    "yButton.on_click(on_button_clicked)\n",
    "lButton = ipywidgets.Button(description=\"L\", tooltip='L')\n",
    "lButton.on_click(on_button_clicked)\n",
    "rButton = ipywidgets.Button(description=\"R\", tooltip='R')\n",
    "rButton.on_click(on_button_clicked)\n",
    "lrButton = ipywidgets.Button(description=\"L+R\", tooltip='2')\n",
    "lrButton.on_click(on_button_clicked)\n",
    "\n",
    "display(ipywidgets.HBox([upButton, downButton, leftButton, rightButton, aButton, bButton, xButton, yButton, lButton, rButton, lrButton]))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Logging data in batch. For testing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import deque\n",
    "logDataSize = 100\n",
    "logDataQueue = deque(maxlen=logDataSize)\n",
    "logDataCounter = 0\n",
    "\n",
    "def logData(d):\n",
    "    global logDataQueue\n",
    "    global logDataCounter\n",
    "    global dataText_w\n",
    "    logDataQueue.append(d)\n",
    "    logDataCounter = logDataCounter + 1\n",
    "    if logDataCounter == logDataSize:\n",
    "        logDataCounter = 0\n",
    "        dataText_w.value = str(logDataQueue)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Major functions for transforming the trt data into game actions. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "#flatten the normalized_peaks from torch.Size([1, 18, 100, 2]) to numpy([1, 36 + 1])\n",
    "#18 peaks * 2 coordinates (y,x) + 1 timestamp\n",
    "def getFlattenPeaks(object_counts, objects, normalized_peaks, timestamp):\n",
    "    count = int(object_counts[0])\n",
    "    flattenPeaks = np.zeros((count,objects[0][0].shape[0]*2 + 1)) #i.e. (1,36 + 1)\n",
    "    for i in range(count):\n",
    "        obj = objects[0][i]\n",
    "        C = obj.shape[0] #18\n",
    "        flattenPeaks[i][C*2] = timestamp #adding timestamp to the end\n",
    "        for j in range(C):\n",
    "            k = int(obj[j]) #obj index, i.e. k=0 for the 1st obj detected\n",
    "            if k >= 0:\n",
    "                    peak = normalized_peaks[0][j][k]\n",
    "                    flattenPeaks[i][j*2] = peak[0]\n",
    "                    flattenPeaks[i][j*2+1] = peak[1]\n",
    "    return flattenPeaks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyPose:\n",
    "    poseIndexNose = 0\n",
    "    poseIndexLeftEye = 1\n",
    "    poseIndexRightEye = 2\n",
    "    poseIndexLeftEar = 3\n",
    "    poseIndexRightEar = 4\n",
    "    poseIndexLeftShoulder = 5\n",
    "    poseIndexRightShoulder = 6\n",
    "    poseIndexLeftElbow = 7\n",
    "    poseIndexRightElbow = 8\n",
    "    poseIndexLeftWrist = 9\n",
    "    poseIndexRightWrist = 10\n",
    "    poseIndexLeftHip = 11\n",
    "    poseIndexRightHip = 12\n",
    "    poseIndexLeftKnee = 13\n",
    "    poseIndexRightKnee = 14\n",
    "    poseIndexLeftAnkle = 15\n",
    "    poseIndexRightAnkle = 16\n",
    "    poseIndexNeck = 17\n",
    "    poseIndexTimeStamp = 18\n",
    "    flattenIndexTimeStamp = poseIndexTimeStamp * 2\n",
    "\n",
    "    def xIndexOf(poseIndex):\n",
    "        return poseIndex * 2 + 1\n",
    "    \n",
    "    def yIndexOf(poseIndex):\n",
    "        return poseIndex * 2\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#return absMaxVal, abs(dVSpeed), startingDVSign, isDVInOneSign\n",
    "def getStatisticsBetweenTwoPeaksInQueue(flattenPeakQueue, flattenIndexA, flattenIndexB):\n",
    "    absMaxVal = 0\n",
    "    lastVal = 0\n",
    "    lastT = 0\n",
    "    dVSum = 0\n",
    "    dTSum = 0\n",
    "    startingDVSign = 0\n",
    "    isDVInOneSign = True\n",
    "    lastAbsVal = 0\n",
    "    isAbsValIncreasing = True\n",
    "    for i, obj in enumerate(flattenPeakQueue):\n",
    "        aVal = obj[flattenIndexA]\n",
    "        bVal = obj[flattenIndexB]\n",
    "        t = obj[MyPose.flattenIndexTimeStamp]\n",
    "        val = bVal - aVal\n",
    "        absVal = abs(val)\n",
    "        if i == 0:\n",
    "            absMaxVal = absVal\n",
    "        else:\n",
    "            if absVal > absMaxVal: absMaxVal = absVal\n",
    "            dV = val - lastVal\n",
    "            dVSign = np.sign(dV)\n",
    "            if startingDVSign == 0:\n",
    "                startingDVSign = dVSign\n",
    "            else:\n",
    "                if dVSign != 0 and startingDVSign != dVSign: isDVInOneSign = False\n",
    "            if absVal < lastAbsVal: isAbsValIncreasing = False\n",
    "            dT = t - lastT\n",
    "            dVSum = dVSum + dV\n",
    "            dTSum = dTSum + dT\n",
    "        lastVal = val\n",
    "        lastT = t\n",
    "        lastAbsVal = absVal\n",
    "    dVSpeed = dVSum / dTSum\n",
    "    return absMaxVal, abs(dVSpeed), startingDVSign, isDVInOneSign, isAbsValIncreasing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#TODO: automatically scale these thresholds according to the full body height (i.e. nose to ankle)\n",
    "isPunchingValThreshold = 0.1\n",
    "isPunchingSpeedThreshold = 0.05\n",
    "\n",
    "def isPunchingWithWrist(flattenPeakQueue, shoulderIndex, wristIndex, valThreshold, speedThreshold):\n",
    "    absMaxVal, absDVSpeed, startingDVSign, isDVInOneSign, isAbsValIncreasing = getStatisticsBetweenTwoPeaksInQueue(flattenPeakQueue, shoulderIndex, wristIndex)\n",
    "    return absMaxVal > valThreshold and absDVSpeed > speedThreshold and isDVInOneSign and isAbsValIncreasing\n",
    "\n",
    "def isPunching(flattenPeakQueue):\n",
    "    global isPunchingValThreshold\n",
    "    global isPunchingSpeedThreshold\n",
    "    return isPunchingWithWrist(flattenPeakQueue, MyPose.xIndexOf(MyPose.poseIndexRightShoulder), MyPose.xIndexOf(MyPose.poseIndexRightWrist), isPunchingValThreshold, isPunchingSpeedThreshold) or isPunchingWithWrist(flattenPeakQueue, MyPose.xIndexOf(MyPose.poseIndexLeftShoulder), MyPose.xIndexOf(MyPose.poseIndexLeftWrist), isPunchingValThreshold, isPunchingSpeedThreshold)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def isPunchingWithWristUpward(flattenPeakQueue, shoulderIndex, wristIndex, valThreshold, speedThreshold):\n",
    "    absMaxVal, absDVSpeed, startingDVSign, isDVInOneSign, isAbsValIncreasing = getStatisticsBetweenTwoPeaksInQueue(flattenPeakQueue, shoulderIndex, wristIndex)\n",
    "    return absMaxVal > valThreshold and absDVSpeed > speedThreshold and isDVInOneSign and isAbsValIncreasing and startingDVSign < 0\n",
    "\n",
    "def isPunchingWithWristDownward(flattenPeakQueue, shoulderIndex, wristIndex, valThreshold, speedThreshold):\n",
    "    absMaxVal, absDVSpeed, startingDVSign, isDVInOneSign, isAbsValIncreasing = getStatisticsBetweenTwoPeaksInQueue(flattenPeakQueue, shoulderIndex, wristIndex)\n",
    "    return absMaxVal > valThreshold and absDVSpeed > speedThreshold and isDVInOneSign and isAbsValIncreasing and startingDVSign > 0\n",
    "\n",
    "def isPunchingUpward(flattenPeakQueue):\n",
    "    global isPunchingValThreshold\n",
    "    global isPunchingSpeedThreshold\n",
    "    return isPunchingWithWristUpward(flattenPeakQueue, MyPose.yIndexOf(MyPose.poseIndexRightShoulder), MyPose.yIndexOf(MyPose.poseIndexRightWrist), isPunchingValThreshold, isPunchingSpeedThreshold) or isPunchingWithWristUpward(flattenPeakQueue, MyPose.yIndexOf(MyPose.poseIndexLeftShoulder), MyPose.yIndexOf(MyPose.poseIndexLeftWrist), isPunchingValThreshold, isPunchingSpeedThreshold)\n",
    "\n",
    "def isPunchingDownward(flattenPeakQueue):\n",
    "    global isPunchingValThreshold\n",
    "    global isPunchingSpeedThreshold\n",
    "    return isPunchingWithWristDownward(flattenPeakQueue, MyPose.yIndexOf(MyPose.poseIndexRightShoulder), MyPose.yIndexOf(MyPose.poseIndexRightWrist), isPunchingValThreshold, isPunchingSpeedThreshold) or isPunchingWithWristDownward(flattenPeakQueue, MyPose.yIndexOf(MyPose.poseIndexLeftShoulder), MyPose.yIndexOf(MyPose.poseIndexLeftWrist), isPunchingValThreshold, isPunchingSpeedThreshold)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "isKickingValThreshold = 0.04\n",
    "isKickingSpeedThreshold = 0.01\n",
    "\n",
    "def isKickingWithAnkle(flattenPeakQueue, hipIndex, ankleIndex, valThreshold, speedThreshold):\n",
    "    absMaxVal, absDVSpeed, startingDVSign, isDVInOneSign, isAbsValIncreasing = getStatisticsBetweenTwoPeaksInQueue(flattenPeakQueue, hipIndex, ankleIndex)\n",
    "    return absMaxVal > valThreshold and absDVSpeed > speedThreshold and isDVInOneSign and isAbsValIncreasing\n",
    "    \n",
    "def isKicking(flattenPeakQueue):\n",
    "    global isKickingValThreshold\n",
    "    global isKickingSpeedThreshold\n",
    "    return isKickingWithAnkle(flattenPeakQueue, MyPose.xIndexOf(MyPose.poseIndexRightHip), MyPose.xIndexOf(MyPose.poseIndexRightAnkle), isKickingValThreshold, isKickingSpeedThreshold) or isKickingWithAnkle(flattenPeakQueue, MyPose.xIndexOf(MyPose.poseIndexLeftHip), MyPose.xIndexOf(MyPose.poseIndexLeftAnkle), isKickingValThreshold, isKickingSpeedThreshold)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "isBlockingValThreshold = 0.02\n",
    "isBlockingSpeedThreshold = 0.01\n",
    "\n",
    "def isBlockingWithWrist(flattenPeakQueue, shoulderIndex, wristIndex, valThreshold, speedThreshold):\n",
    "    absMaxVal, absDVSpeed, startingDVSign, isDVInOneSign, isAbsValIncreasing = getStatisticsBetweenTwoPeaksInQueue(flattenPeakQueue, shoulderIndex, wristIndex)\n",
    "    if absMaxVal < valThreshold and absDVSpeed < speedThreshold:\n",
    "        logData(\"{} {}\".format(absMaxVal, absDVSpeed))\n",
    "    return absMaxVal < valThreshold and absDVSpeed < speedThreshold\n",
    "\n",
    "def isBlocking(flattenPeakQueue):\n",
    "    global isBlockingValThreshold\n",
    "    global isBlockingSpeedThreshold\n",
    "    return isBlockingWithWrist(flattenPeakQueue, MyPose.xIndexOf(MyPose.poseIndexRightShoulder), MyPose.xIndexOf(MyPose.poseIndexRightWrist), isBlockingValThreshold, isBlockingSpeedThreshold) and isBlockingWithWrist(flattenPeakQueue, MyPose.xIndexOf(MyPose.poseIndexLeftShoulder), MyPose.xIndexOf(MyPose.poseIndexLeftWrist), isBlockingValThreshold, isBlockingSpeedThreshold)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def isFlattennPeakQueueValid(flattenPeakQueue):\n",
    "    obj = flattenPeakQueue[0]\n",
    "    xVal = obj[MyPose.xIndexOf(MyPose.poseIndexRightWrist)]\n",
    "    yVal = obj[MyPose.yIndexOf(MyPose.poseIndexRightWrist)]\n",
    "    return xVal != 0 and yVal != 0\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Major game action logic."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#punch \"a\"\n",
    "#kick \"x\"\n",
    "#put down \"y\"\n",
    "#block \"L\"\n",
    "#super \"R\"\n",
    "def getKarateControlFrom(flattenPeakQueue):\n",
    "    if not isFlattennPeakQueueValid(flattenPeakQueue):\n",
    "        return \"N\"\n",
    "    elif isPunching(flattenPeakQueue):\n",
    "        return \"a\"\n",
    "    elif isKicking(flattenPeakQueue):\n",
    "        return \"x\"\n",
    "    elif isPunchingDownward(flattenPeakQueue):\n",
    "        return \"y\"\n",
    "    elif isPunchingUpward(flattenPeakQueue):\n",
    "        return \"R\"    \n",
    "    #elif isBlocking(flattenPeakQueue):\n",
    "    #    return \"L\"    #TODO: Throttle the output of L, as joycontrol take 0.1 sec to emulate key pressing\n",
    "    else:\n",
    "        return \"N\" #N for nothing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lastControlKey = \"a\"\n",
    "\n",
    "def setControlKeyButtonStyle(controlKey, style2Set):\n",
    "    global aButton\n",
    "    global bButton\n",
    "    global xButton\n",
    "    global yButton\n",
    "    global lButton\n",
    "    global rButton\n",
    "    if controlKey == \"a\":\n",
    "        aButton.button_style = style2Set\n",
    "    elif controlKey == \"b\":\n",
    "        bButton.button_style = style2Set\n",
    "    elif controlKey == \"x\":\n",
    "        xButton.button_style = style2Set\n",
    "    elif controlKey == \"y\":\n",
    "        yButton.button_style = style2Set\n",
    "    elif controlKey == \"L\":\n",
    "        lButton.button_style = style2Set\n",
    "    elif controlKey == \"R\":\n",
    "        rButton.button_style = style2Set\n",
    "    \n",
    "def displayAction(controlKey):\n",
    "    global lastControlKey\n",
    "    setControlKeyButtonStyle(lastControlKey, '')\n",
    "    setControlKeyButtonStyle(controlKey, 'success')\n",
    "    lastControlKey = controlKey\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import deque\n",
    "N = 5\n",
    "flattenPeakQueue = deque(maxlen=N)\n",
    "\n",
    "def processTrtData(object_counts, objects, normalized_peaks, timestamp):\n",
    "    global N\n",
    "    global flattenPeakQueue\n",
    "    global dataText_w\n",
    "    global controllerCheckBox_w\n",
    "    global tcpWriter\n",
    "    flattenPeaks = getFlattenPeaks(object_counts, objects, normalized_peaks, timestamp)\n",
    "    if len(flattenPeaks) > 0:\n",
    "        flattenPeak = flattenPeaks[0]\n",
    "        flattenPeakQueue.append(flattenPeak)\n",
    "        if len(flattenPeakQueue) == N:\n",
    "            controlKey = getKarateControlFrom(flattenPeakQueue)\n",
    "            if controllerCheckBox_w.value and controlKey != \"N\":\n",
    "                tcpWriter.write(controlKey.encode())\n",
    "            displayAction(controlKey)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For measuring the response time of the main processing loop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "timeSum = 0\n",
    "timeCount = 0\n",
    "timeCountN = 1000\n",
    "\n",
    "def recordTime(deltaTime):\n",
    "    global timeSum\n",
    "    global timeCount\n",
    "    global dataText_w\n",
    "    global timeCountN\n",
    "    timeSum = timeSum + deltaTime\n",
    "    timeCount = timeCount + 1\n",
    "    if timeCount >= timeCountN:\n",
    "        fps = timeCountN / timeSum\n",
    "        timeSum = 0\n",
    "        timeCount = 0\n",
    "        dataText_w.value = \"fps {}\".format(fps)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we'll define the main execution loop.  This will perform the following steps\n",
    "\n",
    "1.  Preprocess the camera image\n",
    "2.  Execute the neural network\n",
    "3.  Parse the objects from the neural network output\n",
    "4.  Draw the objects onto the camera image\n",
    "5.  Convert the image to JPEG format and stream to the display widget"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def execute(change):\n",
    "    global saveImageCheckBox_w\n",
    "    global saveImageDelayCounter\n",
    "    t0 = time.time()\n",
    "    image = change['new']\n",
    "    data = preprocess(image)\n",
    "    cmap, paf = model_trt(data)\n",
    "    cmap, paf = cmap.detach().cpu(), paf.detach().cpu()\n",
    "    counts, objects, peaks = parse_objects(cmap, paf)#, cmap_threshold=0.15, link_threshold=0.15)\n",
    "    processTrtData(counts, objects, peaks, t0)\n",
    "    draw_objects(image, counts, objects, peaks)\n",
    "    image_w.value = bgr8_to_jpeg(image[:, ::-1, :])\n",
    "    t1 = time.time()\n",
    "    recordTime(t1-t0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we call the cell below it will execute the function once on the current camera frame."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "execute({'new': camera.value})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Call the cell below to attach the execution function to the camera's internal value.  This will cause the execute function to be called whenever a new camera frame is received."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "camera.observe(execute, names='value')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Call the cell below to unattach the camera frame callbacks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "camera.unobserve_all()"
   ]
  },
  {
   "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.8.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}