data_collection.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Classification widget for collision avoidance task"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The widget generated by this notebook facilitates the creation of a dataset suitable for addressing a collision avoidance problem on Jetracer robot. [Jetbot robot code](https://github.com/NVIDIA-AI-IOT/jetbot) is used as a reference for creating the widget.\n",
    "\n",
    "The approach we take to avoiding collisions is to create a virtual \"safety bubble\" around the robot. Within this safety bubble, the robot is able to spin in a circle without hitting any objects. Only a sensor is used: the camera mounted on Jetracer.\n",
    "\n",
    "Two scenarios can be distinguished:\n",
    "- When the robot's \"safety bubble\" is breached. This scenario will be labeled as `blocked`;\n",
    "- When the robot's \"safety bubble\" is not breached. This scenario will be labeled as `free`.\n",
    "\n",
    "> **The steering angle of the JetRacer is limited**. This is an important detail to remember when you're trying to estimate whether the robot's safety bubble is violated or not."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A folder dataset is created and it will contain two sub-folders `free` and `blocked`, where we'll place the images for each scenario. \n",
    "\n",
    "If you refresh the Jupyter file browser on the left, you should now see those directories appear. Next, let's create some buttons that we'll use to save snapshots for each class label. We'll also add some text boxes that will display how many images of each category that we've collected so far. This is useful because we want to make sure we collect about as many free images as blocked images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ipywidgets as widgets\n",
    "from IPython.display import display\n",
    "import cv2\n",
    "import os\n",
    "\n",
    "blocked_dir = 'dataset/blocked'\n",
    "free_dir = 'dataset/free'\n",
    "\n",
    "# we have this \"try/except\" statement because these next functions can throw an error if the directories exist already\n",
    "try:\n",
    "    os.makedirs(free_dir)\n",
    "    os.makedirs(blocked_dir)\n",
    "except FileExistsError:\n",
    "    print('Directories not created because they already exist')\n",
    "    \n",
    "button_layout = widgets.Layout(width='128px', height='64px')\n",
    "free_button = widgets.Button(description='add free', button_style='success', layout=button_layout)\n",
    "blocked_button = widgets.Button(description='add blocked', button_style='danger', layout=button_layout)\n",
    "free_count = widgets.IntText(layout=button_layout, value=len(os.listdir(free_dir)))\n",
    "blocked_count = widgets.IntText(layout=button_layout, value=len(os.listdir(blocked_dir)))\n",
    "\n",
    "#display(widgets.HBox([free_count, free_button]))\n",
    "#display(widgets.HBox([blocked_count, blocked_button]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The camera is initialized."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import traitlets \n",
    "import ipywidgets.widgets as widgets\n",
    "from IPython.display import display\n",
    "from jetcam.csi_camera import CSICamera\n",
    "from jetcam.utils import bgr8_to_jpeg\n",
    "\n",
    "camera = CSICamera(width=224, height=224, capture_fps=65)\n",
    "camera.running = True\n",
    "\n",
    "image_widget = widgets.Image(format='jpeg', width=224, height=224)  # this width and height doesn't necessarily have to match the camera\n",
    "\n",
    "camera_link = traitlets.dlink((camera, 'value'), (image_widget, 'value'), transform=bgr8_to_jpeg)\n",
    "\n",
    "#display(image_widget)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To make sure we don't repeat any file names (even across different machines!) we'll use the `uuid` package in python, which defines the uuid1 method to generate a unique identifier. This unique identifier is generated from information like the current time and the machine address."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from uuid import uuid1\n",
    "\n",
    "count_free = 0\n",
    "count_blocked = 0\n",
    "\n",
    "def save_snapshot(directory):\n",
    "    image_path = os.path.join(directory, str(uuid1()) + '.jpg')\n",
    "    cv2.imwrite(image_path, camera.value.copy())\n",
    "\n",
    "def save_free():\n",
    "    global free_dir, free_count\n",
    "    save_snapshot(free_dir)\n",
    "    free_count.value = len(os.listdir(free_dir))\n",
    "\n",
    "def save_blocked():\n",
    "    global blocked_dir, blocked_count\n",
    "    save_snapshot(blocked_dir)\n",
    "    blocked_count.value = len(os.listdir(blocked_dir))\n",
    "    \n",
    "free_button.on_click(lambda x: save_free())\n",
    "blocked_button.on_click(lambda x: save_blocked())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## All together!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The complete widget is displayed. It presents real-time images from the camera, and the `free` and `blocked` buttons can be used to collect images of the two classes. The number of elements in each class is displayed next to each button."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display(widgets.HBox([free_count, free_button]))\n",
    "display(widgets.HBox([blocked_count, blocked_button]))\n",
    "display(image_widget)"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}