forked from NVIDIA-AI-IOT/jetbot
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge branch 'master' into dev-clickable
- Loading branch information
Showing
13 changed files
with
1,996 additions
and
23 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,313 @@ | ||
{ | ||
"cells": [ | ||
{ | ||
"cell_type": "markdown", | ||
"metadata": {}, | ||
"source": [ | ||
"# Collision Avoidance - Live Demo (Resnet18)\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.resnet18(pretrained=False)\n", | ||
"model.fc = torch.nn.Linear(512, 2)" | ||
] | ||
}, | ||
{ | ||
"cell_type": "markdown", | ||
"metadata": {}, | ||
"source": [ | ||
"Next, load the trained weights from the ``best_model_resnet18.pth`` file that you uploaded" | ||
] | ||
}, | ||
{ | ||
"cell_type": "code", | ||
"execution_count": null, | ||
"metadata": {}, | ||
"outputs": [], | ||
"source": [ | ||
"model.load_state_dict(torch.load('best_model_resnet18.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)\n", | ||
"model = model.eval().half()" | ||
] | ||
}, | ||
{ | ||
"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 doesnt *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 HWC layout to CHW layout\n", | ||
"2. 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", | ||
"3. Transfer the data from CPU memory to GPU memory\n", | ||
"4. Add a batch dimension" | ||
] | ||
}, | ||
{ | ||
"cell_type": "code", | ||
"execution_count": null, | ||
"metadata": {}, | ||
"outputs": [], | ||
"source": [ | ||
"import torchvision.transforms as transforms\n", | ||
"import torch.nn.functional as F\n", | ||
"import cv2\n", | ||
"import PIL.Image\n", | ||
"import numpy as np\n", | ||
"\n", | ||
"mean = torch.Tensor([0.485, 0.456, 0.406]).cuda().half()\n", | ||
"std = torch.Tensor([0.229, 0.224, 0.225]).cuda().half()\n", | ||
"\n", | ||
"normalize = torchvision.transforms.Normalize(mean, std)\n", | ||
"\n", | ||
"def preprocess(image):\n", | ||
" image = PIL.Image.fromarray(image)\n", | ||
" image = transforms.functional.to_tensor(image).to(device).half()\n", | ||
" image.sub_(mean[:, None, None]).div_(std[:, None, None])\n", | ||
" return image[None, ...]" | ||
] | ||
}, | ||
{ | ||
"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." | ||
] | ||
}, | ||
{ | ||
"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", | ||
"\n", | ||
"camera_link = traitlets.dlink((camera, 'value'), (image, 'value'), transform=bgr8_to_jpeg)\n", | ||
"\n", | ||
"display(widgets.HBox([image, blocked_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(0.2)\n", | ||
" else:\n", | ||
" robot.left(0.2)\n", | ||
" \n", | ||
" time.sleep(0.001)\n", | ||
" \n", | ||
"update({'new': camera.value}) # we call the function once to intialize" | ||
] | ||
}, | ||
{ | ||
"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 will move the robot!! Please make sure your robot has clearance. The collision avoidance should work, but the neural\n", | ||
"> network is only as good as the data it's trained on!" | ||
] | ||
}, | ||
{ | ||
"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": [ | ||
"### 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 :)" | ||
] | ||
} | ||
], | ||
"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 | ||
} |
Oops, something went wrong.