dataset_builder.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Camera initialized\n",
      "Controller initialized\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f0dc7165b4d049fea894ebd7c85b08f1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Image(value=b'\\x89PNG\\r\\n\\x1a\\n\\x00\\x00\\x00\\rIHDR\\x00\\x00\\x034\\x00\\x00\\x02h\\x08\\x02\\x00\\x00\\x00?+3\\x18\\x00\\x00…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Shutdown procedure initiated...\n",
      "Shutdown complete\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Exception in thread Thread-4:\n",
      "Traceback (most recent call last):\n",
      "  File \"/usr/lib/python3.6/threading.py\", line 916, in _bootstrap_inner\n",
      "    self.run()\n",
      "  File \"/usr/lib/python3.6/threading.py\", line 864, in run\n",
      "    self._target(*self._args, **self._kwargs)\n",
      "  File \"<ipython-input-2-a0261d72dc1b>\", line 62, in _capture_frames\n",
      "    re, image = self.cap.read()\n",
      "cv2.error: /home/nvidia/build_opencv/opencv/modules/videoio/src/cap_gstreamer.cpp:447: error: (-2) GStreamer: unable to start pipeline\n",
      " in function icvStartPipeline\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import atexit\n",
    "import threading\n",
    "from math import atan2, cos, pi, sin\n",
    "from os import makedirs\n",
    "from time import sleep\n",
    "\n",
    "import cv2\n",
    "import ipywidgets.widgets as widgets\n",
    "import numpy as np\n",
    "import traitlets\n",
    "from jetbot import Robot\n",
    "from traitlets.config.configurable import SingletonConfigurable\n",
    "\n",
    "import xbox as xbox_controller\n",
    "\n",
    "\n",
    "JOYSTICK_DZ_THRESH: float = 0.15\n",
    "RIGHT_TRIGGER_THRESH: float = 0.1\n",
    "PRECISION: int = 3\n",
    "IMAGE_WIDTH: int = 820\n",
    "IMAGE_HEIGHT: int = 616\n",
    "COLLECTION_TPS: int = 30\n",
    "\n",
    "FRAMES_DIRECTORY: str = './test2'\n",
    "FRAMES_START_INDEX: int = 0\n",
    "    \n",
    "def remap_vector(vector):\n",
    "    return vector * 0.5 / 1 + 0.25\n",
    "\n",
    "\n",
    "class Camera(SingletonConfigurable):\n",
    "    \n",
    "    value = traitlets.Any()\n",
    "    \n",
    "    # config\n",
    "    width = traitlets.Integer(default_value=IMAGE_WIDTH).tag(config=True)\n",
    "    height = traitlets.Integer(default_value=IMAGE_HEIGHT).tag(config=True)\n",
    "    fps = traitlets.Integer(default_value=21).tag(config=True)\n",
    "    capture_width = traitlets.Integer(default_value=3280).tag(config=True)\n",
    "    capture_height = traitlets.Integer(default_value=2464).tag(config=True)\n",
    "\n",
    "    def __init__(self, *args, **kwargs):\n",
    "        self.value = np.empty((self.height, self.width, 3), dtype=np.uint8)\n",
    "        super(Camera, self).__init__(*args, **kwargs)\n",
    "\n",
    "        try:\n",
    "            self.cap = cv2.VideoCapture(self._gst_str(), cv2.CAP_GSTREAMER)\n",
    "\n",
    "            re, image = self.cap.read()\n",
    "\n",
    "            if not re:\n",
    "                raise RuntimeError('Could not read image from camera.')\n",
    "\n",
    "            self.value = image\n",
    "            self.start()\n",
    "        except:\n",
    "            self.stop()\n",
    "            raise RuntimeError(\n",
    "                'Could not initialize camera.  Please see error trace.')\n",
    "\n",
    "        atexit.register(self.stop)\n",
    "\n",
    "    def _capture_frames(self):\n",
    "        while True:\n",
    "            re, image = self.cap.read()\n",
    "            if re:\n",
    "                self.value = image\n",
    "            else:\n",
    "                break\n",
    "                \n",
    "    def _gst_str(self):\n",
    "        return 'nvarguscamerasrc ! video/x-raw(memory:NVMM), width=%d, height=%d, format=(string)NV12, framerate=(fraction)%d/1 ! nvvidconv ! video/x-raw, width=(int)%d, height=(int)%d, format=(string)BGRx ! videoconvert ! appsink' % (\n",
    "                self.capture_width, self.capture_height, self.fps, self.width, self.height)\n",
    "    \n",
    "    def start(self):\n",
    "        if not self.cap.isOpened():\n",
    "            self.cap.open(self._gst_str(), cv2.CAP_GSTREAMER)\n",
    "        if not hasattr(self, 'thread') or not self.thread.isAlive():\n",
    "            self.thread = threading.Thread(target=self._capture_frames)\n",
    "            self.thread.start()\n",
    "\n",
    "    def stop(self):\n",
    "        if hasattr(self, 'cap'):\n",
    "            self.cap.release()\n",
    "        if hasattr(self, 'thread'):\n",
    "            self.thread.join()\n",
    "            \n",
    "    def restart(self):\n",
    "        self.stop()\n",
    "        self.start()\n",
    "\n",
    "\n",
    "class MotorValues(traitlets.HasTraits):\n",
    "    motor_speed = traitlets.Float()\n",
    "    \n",
    "    \n",
    "class Collector(traitlets.HasTraits):\n",
    "    def __init__(self, frames_directory: str, frame_index: int):\n",
    "        self.frames_directory = frames_directory\n",
    "        self.frame_index = frame_index\n",
    "        self.should_shutdown = False\n",
    "        self.should_record = False\n",
    "        \n",
    "        self.robot: Robot = Robot()\n",
    "        print('')\n",
    "        self.camera: Camera = Camera.instance()\n",
    "        print('Camera initialized')\n",
    "        self.image_latest = widgets.Image(format='png', width=IMAGE_WIDTH, height=IMAGE_HEIGHT)\n",
    "        \n",
    "        self.controls = {'left_motor': MotorValues(),\n",
    "                         'right_motor': MotorValues()}\n",
    "        self.controller = xbox_controller.Joystick()\n",
    "        print('Controller initialized')\n",
    "        self.ls_x = 0.0\n",
    "        self.ls_y = 0.0\n",
    "        self.right_trigger = 0.0\n",
    "        self.steering_vector = 0.0\n",
    "\n",
    "        # Dynamically link motor controller values to class attribute\n",
    "        self.left_link = traitlets.dlink((self.controls['left_motor'], 'motor_speed'), (self.robot.left_motor, 'value'))\n",
    "        self.right_link = traitlets.dlink((self.controls['right_motor'], 'motor_speed'), (self.robot.right_motor, 'value'))\n",
    "        \n",
    "        # Dynamically link camera output to class attribute\n",
    "        self.camera_link = traitlets.dlink((self.camera, 'value'), (self.image_latest, 'value'), transform=lambda value: bytes(cv2.imencode('.png', value)[1]))\n",
    "        \n",
    "        # Display latest frame from camera\n",
    "        display(self.image_latest)\n",
    "        \n",
    "        self.threads = {'movement': threading.Thread(name='thread_movement', target=self.monitor_status),\n",
    "                       'camera': threading.Thread(name='thread_camera', target=self.record_frames)}\n",
    "        \n",
    "        for thread in self.threads.values():\n",
    "            thread.start()\n",
    "        \n",
    "    def shutdown(self):\n",
    "        print('Shutdown procedure initiated...')\n",
    "        self.controller.close()\n",
    "        self.left_link.unlink()\n",
    "        self.right_link.unlink()\n",
    "        self.camera_link.unlink()\n",
    "        self.camera.stop()\n",
    "        self.robot.stop()\n",
    "        print('Shutdown complete')\n",
    "        exit(0)\n",
    "        \n",
    "    def save_frame(self, steering_vector: float = 0):\n",
    "        \"\"\"Saves a frame into a file in the frames_directory.\"\"\"\n",
    "        file_path = f'{self.frames_directory}/{self.frame_index}-{steering_vector}.png'\n",
    "        with open(file_path, 'wb') as f:\n",
    "            f.write(self.image_latest.value)\n",
    "            \n",
    "    def calculate_steering_vector(self):\n",
    "        return atan2(self.ls_y, self.ls_x) / pi\n",
    "        \n",
    "    def record_frames(self):\n",
    "        while not self.should_shutdown:\n",
    "            # Ignore values of 0 and when y is negative (do not record when reversing)\n",
    "            if self.ls_y > 0 and self.right_trigger > 0:\n",
    "                self.save_frame(self.steering_vector)\n",
    "                self.frame_index += 1\n",
    "            \n",
    "            # Throttle number of frames saved\n",
    "            sleep(1. / COLLECTION_TPS)\n",
    "            \n",
    "    \n",
    "    def monitor_status(self):\n",
    "        while not self.should_shutdown:\n",
    "            # Handle left joystick axes positions: right and up are positive (after inverting y axis)\n",
    "            self.ls_x = round(self.controller.leftX(), PRECISION)\n",
    "            self.ls_y = round(self.controller.leftY(), PRECISION)\n",
    "            \n",
    "             # Handle right trigger position, normalize value\n",
    "            self.right_trigger = round(self.controller.rightTrigger(), PRECISION)\n",
    "\n",
    "            # Compensate for controller deadzone\n",
    "            if abs(self.ls_x) < JOYSTICK_DZ_THRESH:\n",
    "                self.ls_x = 0.0\n",
    "            if abs(self.ls_y) < JOYSTICK_DZ_THRESH:\n",
    "                self.ls_y = 0.0\n",
    "            if self.right_trigger < RIGHT_TRIGGER_THRESH:\n",
    "                self.right_trigger = 0.0\n",
    "            \n",
    "            # Remap steering vector to a smaller range to reduce sharp turns\n",
    "            self.steering_vector = self.calculate_steering_vector()\n",
    "            remap_rads = remap_vector(self.steering_vector) * pi\n",
    "            \n",
    "            # Grab x and y coordinates for remapped vector\n",
    "            x = cos(remap_rads)\n",
    "            y = sin(remap_rads)\n",
    "\n",
    "            # Calculate intermediaries for tank control conversion\n",
    "            v = (1 - abs(x)) * (y / 1) + y\n",
    "            w = (1 - abs(y)) * (x / 1) + x\n",
    "\n",
    "            # Translate intermediaries and apply right trigger speed multiplier to compute tank control values\n",
    "            self.controls['left_motor'].motor_speed = (v + w) / 2 * self.right_trigger\n",
    "            self.controls['right_motor'].motor_speed = (v - w) / 2 * self.right_trigger\n",
    "\n",
    "                \n",
    "if __name__ == '__main__':\n",
    "    makedirs(FRAMES_DIRECTORY, exist_ok=True)\n",
    "    try:\n",
    "        collector = Collector(FRAMES_DIRECTORY, FRAMES_START_INDEX)\n",
    "    except Exception as e:\n",
    "        collector.shutdown()\n",
    "        print(e)\n",
    "        \n",
    "    try:\n",
    "        for thread in collector.threads.values():\n",
    "            thread.join()\n",
    "    except KeyboardInterrupt:\n",
    "        collector.should_shutdown = True\n",
    "        for thread in collector.threads.values():\n",
    "            thread.join()\n",
    "            \n",
    "        collector.shutdown()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'collector' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-1-734b6dfde521>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0;31m# Cell for manual robot shutdown\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m \u001b[0mcollector\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mleft_link\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0munlink\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      3\u001b[0m \u001b[0mcollector\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mright_link\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0munlink\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[0mcollector\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcamera_link\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0munlink\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m \u001b[0mcollector\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcamera\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstop\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mNameError\u001b[0m: name 'collector' is not defined"
     ]
    }
   ],
   "source": [
    "# Cell for manual robot shutdown\n",
    "collector.left_link.unlink()\n",
    "collector.right_link.unlink()\n",
    "collector.camera_link.unlink()\n",
    "collector.camera.stop()\n",
    "collector.controller.close()\n",
    "collector.robot.stop()"
   ]
  },
  {
   "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.8"
  },
  "pycharm": {
   "stem_cell": {
    "cell_type": "raw",
    "metadata": {
     "collapsed": false
    },
    "source": []
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}