autopilot_data_collection.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### IMPORTS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ipywidgets\n",
    "import cv2\n",
    "from IPython.display import display\n",
    "from jetracer.nvidia_racecar import NvidiaRacecar\n",
    "from jetcam.csi_camera import CSICamera\n",
    "import traitlets\n",
    "import time\n",
    "import shutil\n",
    "import os\n",
    "import pandas as pd\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>\n",
    "\n",
    "### HYPERPARAMETERS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "STEERING_OFFSET = 0.035\n",
    "THROTTLE_GAIN = 0.7\n",
    "\n",
    "CAMERA_WIDTH = 448\n",
    "CAMERA_HEIGHT = 336\n",
    "\n",
    "FPS = 10\n",
    "SHOW_CAMERA_PREVIEW = False\n",
    "\n",
    "DATASETS_DIR = \"/home/greg/datasets/\"\n",
    "TMP_DATASET_DIR = DATASETS_DIR + \"tmp/\"\n",
    "ANNOTATIONS_FILE = \"annotations.csv\"\n",
    "TMP_ANNOTATIONS = TMP_DATASET_DIR + ANNOTATIONS_FILE\n",
    "\n",
    "DATASET_MODE = \"training\"\n",
    "DATASET_NAME = \"3\"\n",
    "MAIN_DATASET_DIR = DATASETS_DIR + DATASET_NAME + \"_\" + DATASET_MODE + \"/\"\n",
    "MAIN_ANNOTATIONS = MAIN_DATASET_DIR + ANNOTATIONS_FILE"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>\n",
    "\n",
    "### DATA COLLECTION"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def reset_temp_dataset_dir():\n",
    "    if not os.path.exists(TMP_DATASET_DIR):\n",
    "        os.makedirs(TMP_DATASET_DIR)\n",
    "    else:\n",
    "        shutil.rmtree(TMP_DATASET_DIR)\n",
    "        os.makedirs(TMP_DATASET_DIR)\n",
    "\n",
    "reset_temp_dataset_dir()\n",
    "\n",
    "if not os.path.exists(MAIN_DATASET_DIR):\n",
    "    os.makedirs(MAIN_DATASET_DIR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def start_recording():\n",
    "    reset_temp_dataset_dir()\n",
    "    \n",
    "def save_recording():\n",
    "    for file in os.listdir(TMP_DATASET_DIR):\n",
    "        if file.endswith('.csv'):\n",
    "            if os.path.exists(MAIN_ANNOTATIONS) and os.stat(MAIN_ANNOTATIONS).st_size > 0:\n",
    "                with open(MAIN_ANNOTATIONS, 'a') as main:\n",
    "                    with open(TMP_ANNOTATIONS) as tmp:\n",
    "                        for line in tmp:\n",
    "                            main.write(line)\n",
    "                        tmp.close()\n",
    "                    main.close()\n",
    "                continue\n",
    "        shutil.move(TMP_DATASET_DIR+file, MAIN_DATASET_DIR+file)\n",
    "    \n",
    "    reset_temp_dataset_dir()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>\n",
    "\n",
    "### CAR CONTROLLER"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "car = NvidiaRacecar()\n",
    "car.throttle_gain = THROTTLE_GAIN\n",
    "car.steering_offset = STEERING_OFFSET"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1b465c7e4863499da0c4c69a746760b9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Controller()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "controller = ipywidgets.widgets.Controller(index=0)\n",
    "display(controller)\n",
    "time.sleep(1.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def clamp(value, val_min=-1.0, val_max=1.0):\n",
    "    return min(val_max, max(val_min, value))\n",
    "\n",
    "def is_valid_press(x):\n",
    "    return x['name'] == 'pressed' and x['new']\n",
    "\n",
    "# Specific to Gamepad Mode 2\n",
    "\n",
    "steering_link = traitlets.dlink((controller.axes[2], 'value'), (car, 'steering'), transform=lambda x: clamp(-x))\n",
    "throttle_link = traitlets.dlink((controller.axes[1], 'value'), (car, 'throttle'), transform=lambda x: clamp(x))\n",
    "\n",
    "brake_throttle_link = traitlets.dlink((controller.buttons[7], 'value'), (car, 'throttle'), transform=lambda x: 0.0)\n",
    "half_throttle_link = traitlets.dlink((controller.buttons[5], 'value'), (car, 'throttle'), transform=lambda x: -0.5 if x > 0.5 else 0)\n",
    "\n",
    "controller.buttons[2].observe(lambda x: start_recording() if is_valid_press(x) else ())\n",
    "controller.buttons[1].observe(lambda x: save_recording() if is_valid_press(x) else ())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>\n",
    "\n",
    "### CAMERA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "camera = CSICamera(width=CAMERA_WIDTH, height=CAMERA_HEIGHT, capture_fps=FPS)\n",
    "camera.running = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def save_annotated_camera_frame(frame):\n",
    "    timestamp = str(int(time.time()*1000))\n",
    "    encoded_image = cv2.imencode('.jpg', frame)[1]\n",
    "    \n",
    "    def save_camera_frame():\n",
    "        cv2.imwrite(TMP_DATASET_DIR+timestamp+\".jpg\", frame)\n",
    "        \n",
    "    def save_annotation():\n",
    "        with open(TMP_ANNOTATIONS, 'a') as f: \n",
    "            f.write(timestamp + \", \" + str(round(car.steering, 2)) + \", \" + str(round(car.throttle, 2)) + \"\\n\")\n",
    "            f.close()\n",
    "    \n",
    "    save_camera_frame()\n",
    "    save_annotation()\n",
    "    \n",
    "    encoded_image = cv2.imencode('.jpg',frame)[1]\n",
    "    return bytes(encoded_image)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "image = ipywidgets.Image(format='jpeg')\n",
    "if SHOW_CAMERA_PREVIEW:\n",
    "    display(image)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "camera_link = traitlets.dlink((camera, 'value'), (image, 'value'), transform=save_annotated_camera_frame)"
   ]
  }
 ],
 "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
}