pupbot.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# DGMD S-17 Summer 2020 - Project PupBot\n",
    "\n",
    "## Team Members:\n",
    ">- Chris Crane\n",
    ">- Robert Clapp\n",
    ">- Andrew Pham\n",
    ">- James Sun\n",
    "\n",
    "## Project Objectives:\n",
    "**Enter Project Objectives Here\n",
    "\n",
    "### Notes:\n",
    "####**Do not run the notebook in its entirety.** \n",
    "\n",
    "####**Restart kernel between the running of each section**\n",
    "\n",
    "1. The robot will need to be trained for each member's environment using the **Data Collection** section. \n",
    "2. All new captured images should be saved to the `dataset` folder.\n",
    "3. After finishing data collection, run the **Train Model** section. The trained model should be saved in the file `best_model.pth`\n",
    "4. After finishing training the model, run the **Main Entry Point** section.\n",
    "5. It may be necessary to reboot the robot between data collection and the main entry point to get the camera resource release"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#######################################################################################\n",
    "## Main Entry Point\n",
    "#######################################################################################"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Library Imports\n",
    "import torch\n",
    "import torchvision\n",
    "import torch.nn.functional as F\n",
    "import cv2\n",
    "import numpy as np\n",
    "from IPython.display import display\n",
    "import ipywidgets.widgets as widgets\n",
    "import traitlets\n",
    "from jetbot import Heartbeat, ObjectDetector, Camera, Robot, bgr8_to_jpeg\n",
    "import json\n",
    "import time\n",
    "\n",
    "# Load COCO object detection labels\n",
    "with open('coco_label_map.json', 'r') as f:\n",
    "    coco_labels = json.load(f)\n",
    "f.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = ObjectDetector('ssd_mobilenet_v2_coco.engine')\n",
    "\n",
    "camera = Camera.instance(width=300, height=300, fps=6)\n",
    "\n",
    "robot = Robot()\n",
    "robot.stop() # Sometimes when the Robot instance is created, the robot starting moving\n",
    "\n",
    "heartbeat = Heartbeat()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "collision_model = torchvision.models.alexnet(pretrained=False)\n",
    "collision_model.classifier[6] = torch.nn.Linear(collision_model.classifier[6].in_features, 2)\n",
    "collision_model.load_state_dict(torch.load('best_model.pth'))\n",
    "device = torch.device('cuda')\n",
    "collision_model = collision_model.to(device)\n",
    "\n",
    "mean = 255.0 * np.array([0.485, 0.456, 0.406])\n",
    "stdev = 255.0 * np.array([0.229, 0.224, 0.225])\n",
    "\n",
    "normalize = torchvision.transforms.Normalize(mean, stdev)\n",
    "\n",
    "def preprocess(camera_value):\n",
    "    global device, normalize\n",
    "    x = camera_value\n",
    "    x = cv2.resize(x, (224, 224))\n",
    "    x = cv2.cvtColor(x, cv2.COLOR_BGR2RGB)\n",
    "    x = x.transpose((2, 0, 1))\n",
    "    x = torch.from_numpy(x).float()\n",
    "    x = normalize(x)\n",
    "    x = x.to(device)\n",
    "    x = x[None, ...]\n",
    "    return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random as rnd\n",
    "\n",
    "blocked_widget = widgets.FloatSlider(min=0.0, max=1.0, value=0.0, description='blocked')\n",
    "image_widget = widgets.Image(format='jpeg', width=300, height=300)\n",
    "label_widget = widgets.IntText(value=1, description='tracked label') # Looking for a Person (value = 1)\n",
    "speed_widget = widgets.FloatSlider(value=0.45, min=0.0, max=1.0, step=0.05, description='speed')\n",
    "turn_gain_widget = widgets.FloatSlider(value=0.6, min=0.0, max=2.0, description='turn gain')\n",
    "stop_button = widgets.Button(description='Stop', button_style='danger')\n",
    "start_button = widgets.Button(description='Restart', button_style='success')\n",
    "det_center_text = widgets.Text(value='')\n",
    "det_confidence_text = widgets.Text(value='')\n",
    "det_label_text = widgets.Text(value='')\n",
    "dist_text = widgets.Text(value='')\n",
    "\n",
    "display(widgets.VBox([\n",
    "    widgets.HBox([image_widget, widgets.VBox([blocked_widget, widgets.HBox([stop_button, start_button])])]),\n",
    "    widgets.HBox([label_widget, det_center_text, dist_text]),\n",
    "    widgets.HBox([speed_widget, det_confidence_text]),\n",
    "    widgets.HBox([turn_gain_widget, det_label_text])\n",
    "]))\n",
    "\n",
    "width = int(image_widget.width)\n",
    "height = int(image_widget.height)\n",
    "\n",
    "#Returns the COCO label given label id\n",
    "def get_coco_label(label_id):\n",
    "    label = 'UNKNOWN'\n",
    "    for l in coco_labels:\n",
    "        if l['id'] == label_id:\n",
    "            label = l['label']\n",
    "            break\n",
    "    return label\n",
    "\n",
    "#Turns robot 90 degrees\n",
    "def turn_90_degrees():\n",
    "    robot.left(float(speed_widget.value))\n",
    "    time.sleep(0.50)\n",
    "\n",
    "#Determine if the detection found is one being looked for and meets the confidence threshold\n",
    "def is_matching_detection(detection):\n",
    "    return detection['label'] == int(label_widget.value) and detection['confidence'] > 0.70 \n",
    "    \n",
    "#Computes the center x, y coordinates of the object\n",
    "def detection_center(detection):\n",
    "    bbox = detection['bbox']\n",
    "    center_x = (bbox[0] + bbox[2]) / 2.0 - 0.5\n",
    "    center_y = (bbox[1] + bbox[3]) / 2.0 - 0.5\n",
    "    return (center_x, center_y)\n",
    "\n",
    "#Calculates the distance to detection\n",
    "def distance_to_detection(detection):\n",
    "    (x1, y1, x2, y2) = detection['bbox']\n",
    "    fl = 1.85 #focal length mm\n",
    "    sh = 75 #sensor height mm\n",
    "    ph = 1700 #average human height mm\n",
    "    ih = height #image height pixels\n",
    "    oh = (y2 - y1) * height #object height in pixels\n",
    "    return (fl * ph * height) / (oh * sh)\n",
    "\n",
    "#Determines if the robot is too close to the detected object\n",
    "def is_too_close(detection):\n",
    "    return distance_to_detection(detection) < 100\n",
    "\n",
    "#Moves the robot away from the detected object\n",
    "def run_away():\n",
    "    det_label_text.value = det_label_text.value + ': Running Away!!'\n",
    "    turn_90_degrees()\n",
    "    robot.forward(float(speed_widget.value * 1.5))\n",
    "\n",
    "#Computes the length of the 2D vector\n",
    "def norm(vec):\n",
    "    return np.sqrt(vec[0]**2 + vec[1]**2)\n",
    "\n",
    "#Finds the detection closest to the image center\n",
    "def closest_detection(detections):\n",
    "    closest_detection = None\n",
    "    for det in detections:\n",
    "        center = detection_center(det)\n",
    "        if closest_detection is None:\n",
    "            closest_detection = det\n",
    "        elif norm(detection_center(det)) < norm(detection_center(closest_detection)):\n",
    "            closest_detection = det\n",
    "    return closest_detection\n",
    "\n",
    "#Process camera input, find detections, follow target if identidied\n",
    "def execute(change):\n",
    "    image = change['new']\n",
    "    \n",
    "    # execute collision model to determine if blocked\n",
    "    collision_output = collision_model(preprocess(image)).detach().cpu()\n",
    "    prob_blocked = float(F.softmax(collision_output.flatten(), dim=0)[0])\n",
    "    blocked_widget.value = prob_blocked\n",
    "    \n",
    "    # turn left if blocked\n",
    "    if prob_blocked > 0.5:\n",
    "        robot.left(float(speed_widget.value))\n",
    "        image_widget.value = bgr8_to_jpeg(image)\n",
    "        return\n",
    "        \n",
    "    # compute all detected objects\n",
    "    detections = model(image)\n",
    "    \n",
    "    # draw all detections on image\n",
    "    for det in detections[0]:\n",
    "        if det['confidence'] > 0.5:\n",
    "            (x, y, w, h) = det['bbox']\n",
    "            l = get_coco_label(det['label'])\n",
    "            cv2.rectangle(image, (int(width * x), int(width * y)), (int(width * w), int(height * h)), (255, 0, 0), 2)\n",
    "            cv2.putText(image, l, (int(width * x), int(height * y)), cv2.FONT_HERSHEY_SIMPLEX, 1, (255,0,0), 2, cv2.LINE_AA)\n",
    "    \n",
    "    # get detection closest to center of field of view that match selected class label and draw it\n",
    "    det = closest_detection([d for d in detections[0] if is_matching_detection(d)])\n",
    "    \n",
    "    if det is None:\n",
    "        # otherwise go forward if no target detected\n",
    "        robot.forward(float(speed_widget.value))\n",
    "    else:\n",
    "        (x, y, w, h) = det['bbox']\n",
    "        l = get_coco_label(det['label'])\n",
    "        cv2.rectangle(image, (int(width * x), int(height * y)), (int(width * w), int(height * h)), (0, 255, 0), 2)\n",
    "        cv2.putText(image, l, (int(width * x), int(height * y)), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2, cv2.LINE_AA)\n",
    "        # move robot forward and steer proportional target's x-distance from center\n",
    "        center = detection_center(det)\n",
    "        det_center_text.value = str(center)\n",
    "        det_confidence_text.value = str(det['confidence'])\n",
    "        det_label_text.value = l\n",
    "        dist_text.value = str(distance_to_detection(det))\n",
    "        if is_too_close(det) == True:\n",
    "            run_away()\n",
    "        else:\n",
    "            robot.set_motors(\n",
    "                float(speed_widget.value + turn_gain_widget.value * center[0]),\n",
    "                float(speed_widget.value - turn_gain_widget.value * center[0])\n",
    "            )\n",
    "    \n",
    "    # update image widget\n",
    "    image_widget.value = bgr8_to_jpeg(image)\n",
    "    \n",
    "def stop(change):\n",
    "    camera.unobserve_all()\n",
    "    time.sleep(0.5)\n",
    "    robot.stop()\n",
    "\n",
    "def restart(change):\n",
    "    camera.unobserve_all()\n",
    "    camera.observe(execute, names='value')\n",
    "\n",
    "# this function will be called when heartbeat 'alive' status changes\n",
    "def handle_heartbeat_status(change):\n",
    "    if change['new'] == Heartbeat.Status.dead:\n",
    "        #stop(change)\n",
    "        det_center_text.value = 'Lost Connection'\n",
    "        \n",
    "stop_button.on_click(stop)\n",
    "start_button.on_click(restart)\n",
    "\n",
    "execute({'new': camera.value})\n",
    "camera.unobserve_all()\n",
    "camera.observe(execute, names='value')\n",
    "heartbeat.observe(handle_heartbeat_status, names='status')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##############################################################################\n",
    "## Data Collection\n",
    "##############################################################################"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "# Display live camera feed\n",
    "import traitlets\n",
    "import ipywidgets.widgets as widgets\n",
    "from IPython.display import display\n",
    "from jetbot import Camera, bgr8_to_jpeg\n",
    "from uuid import uuid1\n",
    "import os\n",
    "\n",
    "camera = Camera.instance(width=224, height=224)\n",
    "\n",
    "image = 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, 'value'), transform=bgr8_to_jpeg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "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 becasue they already exist')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "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)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "def save_snapshot(directory):\n",
    "    image_path = os.path.join(directory, str(uuid1()) + '.jpg')\n",
    "    with open(image_path, 'wb') as f:\n",
    "        f.write(image.value)\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())\n",
    "\n",
    "display(image)\n",
    "display(widgets.HBox([free_count, free_button]))\n",
    "display(widgets.HBox([blocked_count, blocked_button]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#######################################################################################\n",
    "## Train Model\n",
    "#######################################################################################"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.optim as optim\n",
    "import torch.nn.functional as F\n",
    "import torchvision\n",
    "import torchvision.datasets as datasets\n",
    "import torchvision.models as models\n",
    "import torchvision.transforms as transforms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "# Create dataset instance\n",
    "dataset = datasets.ImageFolder(\n",
    "    'dataset',\n",
    "    transforms.Compose([\n",
    "        transforms.ColorJitter(0.1, 0.1, 0.1, 0.1),\n",
    "        transforms.Resize((224, 224)),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n",
    "    ])\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "# Split dataset into train and test sets\n",
    "train_dataset, test_dataset = torch.utils.data.random_split(dataset, [len(dataset) - 50, 50])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "# Create data loaders to load data in batches\n",
    "train_loader = torch.utils.data.DataLoader(\n",
    "    train_dataset,\n",
    "    batch_size=16,\n",
    "    shuffle=True,\n",
    "    num_workers=4\n",
    ")\n",
    "\n",
    "test_loader = torch.utils.data.DataLoader(\n",
    "    test_dataset,\n",
    "    batch_size=16,\n",
    "    shuffle=True,\n",
    "    num_workers=4\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "# Define the neural network\n",
    "model = models.alexnet(pretrained=True)\n",
    "model.classifier[6] = torch.nn.Linear(model.classifier[6].in_features, 2)\n",
    "device = torch.device('cuda')\n",
    "model = model.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "# Train the neural network\n",
    "NUM_EPOCHS = 30\n",
    "BEST_MODEL_PATH = 'best_model.pth'\n",
    "best_accuracy = 0.0\n",
    "\n",
    "optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)\n",
    "\n",
    "for epoch in range(NUM_EPOCHS):\n",
    "    \n",
    "    for images, labels in iter(train_loader):\n",
    "        images = images.to(device)\n",
    "        labels = labels.to(device)\n",
    "        optimizer.zero_grad()\n",
    "        outputs = model(images)\n",
    "        loss = F.cross_entropy(outputs, labels)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "    \n",
    "    test_error_count = 0.0\n",
    "    for images, labels in iter(test_loader):\n",
    "        images = images.to(device)\n",
    "        labels = labels.to(device)\n",
    "        outputs = model(images)\n",
    "        test_error_count += float(torch.sum(torch.abs(labels - outputs.argmax(1))))\n",
    "    \n",
    "    test_accuracy = 1.0 - float(test_error_count) / float(len(test_dataset))\n",
    "    print('%d: %f' % (epoch, test_accuracy))\n",
    "    if test_accuracy > best_accuracy:\n",
    "        torch.save(model.state_dict(), BEST_MODEL_PATH)\n",
    "        best_accuracy = test_accuracy"
   ]
  }
 ],
 "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.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}