notebooks/collision_avoidance/live_demo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Collision Avoidance - Live Demo\n",
    "\n",
    "In this notebook we'll use the model we trained to detect whether the robot is ``free`` or ``blocked`` to enable a collision avoidance behavior on the robot.  \n",
    "\n",
    "## Load the trained model\n",
    "\n",
    "We'll assumed that you've already downloaded the ``best_model.pth`` to your workstation as instructed in the training notebook.  Now, you should upload this model into this notebook's\n",
    "directory by using the Jupyter Lab upload tool.  Once that's finished there should be a file named ``best_model.pth`` in this notebook's directory.  \n",
    "\n",
    "> Please make sure the file has uploaded fully before calling the next cell\n",
    "\n",
    "Execute the code below to initialize the PyTorch model.  This should look very familiar from the training notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torchvision\n",
    "\n",
    "model = torchvision.models.alexnet(pretrained=False)\n",
    "model.classifier[6] = torch.nn.Linear(model.classifier[6].in_features, 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, load the trained weights from the ``best_model.pth`` file that you uploaded"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.load_state_dict(torch.load('best_model.pth'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Currently, the model weights are located on the CPU memory execute the code below to transfer to the GPU device."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device('cuda')\n",
    "model = model.to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create the preprocessing function\n",
    "\n",
    "We have now loaded our model, but there's a slight issue.  The format that we trained our model doesn't *exactly* match the format of the camera.  To do that, \n",
    "we need to do some *preprocessing*.  This involves the following steps\n",
    "\n",
    "1. Convert from BGR to RGB\n",
    "2. Convert from HWC layout to CHW layout\n",
    "3. Normalize using same parameters as we did during training (our camera provides values in [0, 255] range and training loaded images in [0, 1] range so we need to scale by 255.0\n",
    "4. Transfer the data from CPU memory to GPU memory\n",
    "5. Add a batch dimension"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "import numpy as np\n",
    "\n",
    "mean = 255.0 * np.array([0.485, 0.456, 0.406])\n",
    "stdev = 255.0 * np.array([0.229, 0.224, 0.225])\n",
    "\n",
    "normalize = torchvision.transforms.Normalize(mean, stdev)\n",
    "\n",
    "def preprocess(camera_value):\n",
    "    global device, normalize\n",
    "    x = camera_value\n",
    "    x = cv2.cvtColor(x, cv2.COLOR_BGR2RGB)\n",
    "    x = x.transpose((2, 0, 1))\n",
    "    x = torch.from_numpy(x).float()\n",
    "    x = normalize(x)\n",
    "    x = x.to(device)\n",
    "    x = x[None, ...]\n",
    "    return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Great! We've now defined our pre-processing function which can convert images from the camera format to the neural network input format.\n",
    "\n",
    "Now, let's start and display our camera.  You should be pretty familiar with this by now.  We'll also create a slider that will display the\n",
    "probability that the robot is blocked.  We'll also display a slider that allows us to control the robot's base speed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import traitlets\n",
    "from IPython.display import display\n",
    "import ipywidgets.widgets as widgets\n",
    "from jetbot import Camera, bgr8_to_jpeg\n",
    "\n",
    "camera = Camera.instance(width=224, height=224)\n",
    "image = widgets.Image(format='jpeg', width=224, height=224)\n",
    "blocked_slider = widgets.FloatSlider(description='blocked', min=0.0, max=1.0, orientation='vertical')\n",
    "speed_slider = widgets.FloatSlider(description='speed', min=0.0, max=0.5, value=0.0, step=0.01, orientation='horizontal')\n",
    "\n",
    "camera_link = traitlets.dlink((camera, 'value'), (image, 'value'), transform=bgr8_to_jpeg)\n",
    "\n",
    "display(widgets.VBox([widgets.HBox([image, blocked_slider]), speed_slider]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll also create our robot instance which we'll need to drive the motors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from jetbot import Robot\n",
    "\n",
    "robot = Robot()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we'll create a function that will get called whenever the camera's value changes.  This function will do the following steps\n",
    "\n",
    "1. Pre-process the camera image\n",
    "2. Execute the neural network\n",
    "3. While the neural network output indicates we're blocked, we'll turn left, otherwise we go forward."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import torch.nn.functional as F\n",
    "import time\n",
    "\n",
    "def update(change):\n",
    "    global blocked_slider, robot\n",
    "    x = change['new'] \n",
    "    x = preprocess(x)\n",
    "    y = model(x)\n",
    "    \n",
    "    # we apply the `softmax` function to normalize the output vector so it sums to 1 (which makes it a probability distribution)\n",
    "    y = F.softmax(y, dim=1)\n",
    "    \n",
    "    prob_blocked = float(y.flatten()[0])\n",
    "    \n",
    "    blocked_slider.value = prob_blocked\n",
    "    \n",
    "    if prob_blocked < 0.5:\n",
    "        robot.forward(speed_slider.value)\n",
    "    else:\n",
    "        robot.left(speed_slider.value)\n",
    "    \n",
    "    time.sleep(0.001)\n",
    "        \n",
    "update({'new': camera.value})  # we call the function once to initialize"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Cool! We've created our neural network execution function, but now we need to attach it to the camera for processing. \n",
    "\n",
    "We accomplish that with the ``observe`` function.\n",
    "\n",
    "> WARNING: This code may move the robot!! Adjust the speed slider we defined earlier to control the base robot speed.  Some kits can move fast, so start slow, and gradually increase the value."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "camera.observe(update, names='value')  # this attaches the 'update' function to the 'value' traitlet of our camera"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Awesome! If your robot is plugged in it should now be generating new commands with each new camera frame.  Perhaps start by placing your robot on the ground and seeing what it does when it reaches an obstacle.\n",
    "\n",
    "If you want to stop this behavior, you can unattach this callback by executing the code below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "camera.unobserve(update, names='value')\n",
    "\n",
    "time.sleep(0.1)  # add a small sleep to make sure frames have finished processing\n",
    "\n",
    "robot.stop()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Perhaps you want the robot to run without streaming video to the browser.  You can unlink the camera as below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "camera_link.unlink()  # don't stream to browser (will still run camera)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To continue streaming call the following."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "camera_link.link()  # stream to browser (wont run camera)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Again, let's close the camera conneciton properly so that we can use the camera in other notebooks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "camera.stop()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Conclusion\n",
    "\n",
    "That's it for this live demo!  Hopefully you had some fun and your robot avoided collisions intelligently! \n",
    "\n",
    "If your robot wasn't avoiding collisions very well, try to spot where it fails.  The beauty is that we can collect more data for these failure scenarios\n",
    "and the robot should get even better :)"
   ]
  },
  {
   "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.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}