model_drive.ipynb

{
 "nbformat": 4,
 "nbformat_minor": 2,
 "metadata": {
  "language_info": {
   "name": "python",
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   }
  },
  "orig_nbformat": 2,
  "file_extension": ".py",
  "mimetype": "text/x-python",
  "name": "python",
  "npconvert_exporter": "python",
  "pygments_lexer": "ipython3",
  "version": 3
 },
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import traitlets\n",
    "from math import abs, pi\n",
    "from jetbot import Robot\n",
    "from dataset_builder import Camera\n",
    "from data_preprocess import process_image_for_model\n",
    "\n",
    "\n",
    "def denormalize_vector(vector: float) -> float:\n",
    "    \"\"\"\n",
    "    Denormalize the normalized vector angle so the robot can use it again.\n",
    "    :param vector: normalized vector value between 0.0 and 1.0\n",
    "    :return: vector value between 0.25 and 0.75\n",
    "    \"\"\"\n",
    "    return (vector * 0.5) + 0.25\n",
    "\n",
    "\n",
    "class AutoPilot(traitlets.HasTraits):\n",
    "    def __init__(self):\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",
    "        self.controls = {'left_motor': MotorValues(),\n",
    "                         'right_motor': MotorValues()}\n",
    "\n",
    "        self.driving_model: tf.keras.models.Model = tf.keras.models.load_model('driving_model.model')\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",
    "    def shutdown(self):\n",
    "        print('Shutdown procedure initiated...')\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 drive(self):\n",
    "        while True:\n",
    "            prediction = self.driving_model.predict([process_image_for_model(self.image_latest.value)])\n",
    "            steering_rads = denormalize_vector(prediction[0]) * pi\n",
    "\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 * 0.5\n",
    "            self.controls['right_motor'].motor_speed = (v - w) / 2 * 0.5\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    try:\n",
    "        auto_pilot = AutoPilot()\n",
    "    except Exception as e:\n",
    "        auto_pilot.shutdown()\n",
    "        print(e)\n",
    "    \n",
    "    auto_pilot.drive()\n",
    "    auto_pilot.shutdown()"
   ]
  }
 ]
}