jetbot_traffic_sign_detection.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "exciting-catering",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Imports\n",
    "# Must restart kernel for changes to show up once imported! RClick -> Restart Kernel before running.\n",
    "import robot_control as lfrobot\n",
    "from sign_classifier import SignClassifier\n",
    "from jetbot import Camera\n",
    "from IPython.display import display\n",
    "import ipywidgets\n",
    "import ipywidgets.widgets as widgets\n",
    "import os\n",
    "import time\n",
    "import numpy as np\n",
    "import cv2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "simplified-worry",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialise camera\n",
    "camera = Camera()\n",
    "camera.start()\n",
    "\n",
    "# Define Util functions\n",
    "def crop_img(img, cropx, cropy):\n",
    "    \"\"\"Returns a cropped image with dimensions (cropx, cropy) centred on the input image.\"\"\"\n",
    "    y,x, _ = img.shape\n",
    "    startx = x//2-(cropx//2)\n",
    "    starty = y//2-(cropy//2)    \n",
    "    return img[starty:starty+cropy,startx:startx+cropx, :]\n",
    "\n",
    "def bgr8_to_png(image):\n",
    "    \"\"\"Converts brg8 image into png\"\"\"\n",
    "    return cv2.imencode('.png', image)[1].tobytes()\n",
    "\n",
    "def png_to_bgr8(image):\n",
    "    \"\"\"Converts png image into bgr8\"\"\"\n",
    "    return cv2.imdecode(np.frombuffer(image, dtype=np.uint8) , cv2.IMREAD_UNCHANGED)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "duplicate-lesson",
   "metadata": {},
   "outputs": [],
   "source": [
    "# DEFINE ROBOT BEHAVIOURAL LOGIC\n",
    "\n",
    "# Define constants\n",
    "CONFIDENCE_THRESHOLD = 0.3  # prediction confidence threshold for it to be considered in the window\n",
    "WINDOW_THRESHOLD = 3  # Number of 'confident' predictions the model must make in a row before we take corresponding action \n",
    "\n",
    "SPEED_HIGH = 0.13\n",
    "TURN_SPEED_HIGH = 0.2\n",
    "\n",
    "SPEED_NORMAL = 0.1\n",
    "TURN_SPEED_NORMAL = 0.15\n",
    "\n",
    "DEAD_RECKONING_DURATION = 1  # Seconds\n",
    "STOP_DURATION = 2\n",
    "\n",
    "# Input image dimensions (centred)\n",
    "CROPX = 85\n",
    "CROPY = 85\n",
    "\n",
    "# Define widgets here (helpful for debugging)\n",
    "detections_widget = widgets.Textarea()\n",
    "snapshot_widget = ipywidgets.Image(width=CROPX, height=CROPY)\n",
    "\n",
    "reporting_widget = ipywidgets.VBox([\n",
    "    snapshot_widget,\n",
    "    detections_widget\n",
    "])\n",
    "\n",
    "def execute(image, frame_no, detection_window):\n",
    "    \"\"\"Makes predictions and executes jetbot behaviour instructions accordingly\n",
    "    \n",
    "    We employ both prediction confidence thresholding as well as a sliding window of consecutive predictions.\n",
    "\n",
    "    The model prediction must first have a probability greater than our threshold to be considered.\n",
    "    If a prediction does not meet this threshold, the last window value will be removed.\n",
    "    This aims to penalise instability in model predictions, but not so much that it slows reaction time by too much.\n",
    "\n",
    "    We also employ a sliding window of consecutive predictions so that we can guarantee greater confidence than what can\n",
    "    be decerned only with class probabilities given by the model.\n",
    "    \n",
    "    \"\"\"\n",
    "    \n",
    "    # get prediction\n",
    "    pred = clf.get_prediction(image)[0]\n",
    "    pred_max = np.argmax(pred)\n",
    "\n",
    "    # Update debug widgets\n",
    "    snapshot_widget.value = bgr8_to_png(crop_img(image, CROPX, CROPY))\n",
    "    detections_widget.value = f\"PREDICTION: {pred_max} |{str(pred)}| SEQ {frame_no} | {detection_window}\"\n",
    "    \n",
    "    # Enforce confidence thresholding + sliding window thresholding \n",
    "    if max(pred) > CONFIDENCE_THRESHOLD:\n",
    "        if len(detection_window) == 0 or detection_window[-1] == pred_max:\n",
    "            # append prediction to sliding window\n",
    "            detection_window.append(pred_max)\n",
    "            \n",
    "            # execute instruction if number of consecutive predictions meets threshold\n",
    "            if len(detection_window) > WINDOW_THRESHOLD:\n",
    "#               ======= SPEED UP =======\n",
    "                if pred_max == 0:\n",
    "                    # change to higher speed\n",
    "                    lfrobot.lfSpeed(SPEED_HIGH)\n",
    "                    lfrobot.lfTurnSpeed(TURN_SPEED_HIGH)\n",
    "                    detections_widget.value = detections_widget.value + \" | 55 SPEEDING!!!!!!!!!!!!!\"\n",
    "                    \n",
    "                    # pause line following for 1s\n",
    "                    lfrobot.lfLineFollow(False)\n",
    "                    time.sleep(DEAD_RECKONING_DURATION)\n",
    "                    lfrobot.lfLineFollow(True)\n",
    "                    \n",
    "#               ========= STOP ==========\n",
    "                elif pred_max == 1:\n",
    "                    # stop robot for 2 s\n",
    "                    detections_widget.value = detections_widget.value + \" | STOPPING\"\n",
    "                    lfrobot.lfStop()\n",
    "                    time.sleep(STOP_DURATION)\n",
    "                    lfrobot.lfStart()\n",
    "                    \n",
    "                    # pause predictions and line follwowing for 1s\n",
    "                    detection_window.clear()\n",
    "                    lfrobot.lfLineFollow(False)\n",
    "                    time.sleep(DEAD_RECKONING_DURATION)\n",
    "                    lfrobot.lfLineFollow(True)\n",
    "                    \n",
    "#               ======= RED LIGHT =======\n",
    "                elif pred_max == 2:\n",
    "                    # stop robot motion until green light is seen\n",
    "                    lfrobot.lfStop()\n",
    "                    detections_widget.value = detections_widget.value + \" | RED LIGHT\"\n",
    "                \n",
    "#               ====== GREEN LIGHT ======\n",
    "                elif pred_max == 3:\n",
    "                    # start robot motion\n",
    "                    detections_widget.value = detections_widget.value + \" | GREEN LIGHT\"\n",
    "                    lfrobot.lfStart()\n",
    "                \n",
    "#               ====== YELLOW LIGHT =====\n",
    "                elif pred_max == 4:\n",
    "                    # return robot to original (slower) speed\n",
    "                    detections_widget.value = detections_widget.value + \" | YELLOW LIGHT\"\n",
    "                    lfrobot.lfSpeed(SPEED_NORMAL)\n",
    "                    lfrobot.lfTurnSpeed(TURN_SPEED_NORMAL)\n",
    "                    \n",
    "                detection_window.clear()\n",
    "        else:\n",
    "            # clear prediction window if previous prediction does not match current\n",
    "            detection_window.clear()\n",
    "            \n",
    "    elif len(detection_window) > 0:\n",
    "        # pop last prediction from window if threshold is not reached by any class\n",
    "        detection_window.pop()\n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "sought-cherry",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define constants and model\n",
    "HZ = 60  # Max HZ. It is not a minimum! \n",
    "clf = SignClassifier(\"cmyk-flat-optimised.joblib\")\n",
    "\n",
    "display(reporting_widget)\n",
    "\n",
    "def run_detection():\n",
    "    \"\"\"Runs detections at a given frequency\"\"\"\n",
    "    counter = 0\n",
    "    detection_window = []\n",
    "    # Continuously predict and execute instructions based on predictions\n",
    "    while True:\n",
    "        time.sleep(1/HZ)\n",
    "        execute(crop_img(camera.value, CROPX, CROPY), counter, detection_window)\n",
    "        counter += 1\n",
    "        \n",
    "def save_image(run_label, save_dir, frame_no):\n",
    "    \"\"\"Saves cropped png image file to designated directory\"\"\"\n",
    "    \n",
    "    # Perform crop on current camera value and display on widget\n",
    "    image = crop_img(camera.value, CROPX, CROPY)\n",
    "    time.sleep(1/HZ)\n",
    "    snapshot_widget.value = bgr8_to_png(image)  # Save as png instead of JPEG for lossless data\n",
    "    detections_widget.value = f\"IMAGE NO.: {frame_no}\"\n",
    "    \n",
    "    # Save image as png\n",
    "    filename = f\"{run_label}-{frame_no}.png\"\n",
    "    with open(os.path.join(save_dir, filename), 'wb') as f:\n",
    "        f.write(bgr8_to_png(image))\n",
    "        \n",
    "def collect_loop_data(run_label, save_root, max_count=None):\n",
    "    \"\"\"Collects series of sample images to designated directory\"\"\"\n",
    "    \n",
    "    # Validate save directory\n",
    "    if not os.path.isdir(save_root):\n",
    "        raise RuntimeError(\"Not a valid save root directory!\")\n",
    "    \n",
    "    # Create new directory if it does not exist already\n",
    "    save_dir = os.path.join(save_root, run_label)\n",
    "    if not os.path.exists(save_dir):\n",
    "        os.mkdir(save_dir)\n",
    "    \n",
    "    # Save images\n",
    "    if max_count is None:\n",
    "        frameCount = 0\n",
    "        while True:\n",
    "            save_image(run_label, save_dir, frameCount)\n",
    "            frameCount +=1\n",
    "            \n",
    "    else:\n",
    "        for i in range(max_count):\n",
    "            save_image(run_label, save_dir, i)\n",
    "            \n",
    "# ====== DATA COLLECTION ======\n",
    "# collect_loop_data(\"85-new\", \"./data\", max_count=None)\n",
    "\n",
    "# ===== DETECTION TEST =====\n",
    "# print(cv2.cvtColor(camera.value, cv2.COLOR_BGR2RGB))\n",
    "# np.histogram(camera.value[:, :, 2], bins=256, range=(0, 256))\n",
    "# np.histogram(cv2.cvtColor(camera.value, cv2.COLOR_BGR2RGB)[:, :, 2], bins=256, range=(0, 256))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "otherwise-alfred",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ====== RUN DETECTION ======\n",
    "\n",
    "# Initialise robot\n",
    "lfrobot.lfInit()\n",
    "lfrobot.lfSpeed(SPEED_NORMAL)\n",
    "lfrobot.lfTurnSpeed(TURN_SPEED_NORMAL)\n",
    "lfrobot.lfStart()\n",
    "\n",
    "# Start detection loop\n",
    "run_detection()"
   ]
  }
 ],
 "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": 5
}