PyTorch_Object_Detection.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from models import *\n",
    "from utils import *\n",
    "\n",
    "import os, sys, time, datetime, random\n",
    "import torch\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision import datasets, transforms\n",
    "from torch.autograd import Variable\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.patches as patches\n",
    "from PIL import Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/shubham-sakha/anaconda3/envs/car_env/lib/python3.7/site-packages/torch/nn/_reduction.py:43: UserWarning: size_average and reduce args will be deprecated, please use reduction='mean' instead.\n",
      "  warnings.warn(warning.format(ret))\n"
     ]
    }
   ],
   "source": [
    "config_path='config/yolov3.cfg'\n",
    "weights_path='config/yolov3.weights'\n",
    "class_path='config/coco.names'\n",
    "img_size=416\n",
    "conf_thres=0.8\n",
    "nms_thres=0.4\n",
    "\n",
    "# Load model and weights\n",
    "model = Darknet(config_path, img_size=img_size)\n",
    "model.load_weights(weights_path)\n",
    "#model.cuda()\n",
    "model.eval()\n",
    "classes = utils.load_classes(class_path)\n",
    "Tensor = torch.cuda.FloatTensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def detect_image(img):\n",
    "    # scale and pad image\n",
    "    ratio = min(img_size/img.size[0], img_size/img.size[1])\n",
    "    imw = round(img.size[0] * ratio)\n",
    "    imh = round(img.size[1] * ratio)\n",
    "    img_transforms = transforms.Compose([ transforms.Resize((imh, imw)),\n",
    "         transforms.Pad((max(int((imh-imw)/2),0), max(int((imw-imh)/2),0), max(int((imh-imw)/2),0), max(int((imw-imh)/2),0)),\n",
    "                        (128,128,128)),\n",
    "         transforms.ToTensor(),\n",
    "         ])\n",
    "    # convert image to Tensor\n",
    "    image_tensor = img_transforms(img).float()\n",
    "    image_tensor = image_tensor.unsqueeze_(0)\n",
    "    input_img = Variable(image_tensor.type(Tensor))\n",
    "    # run inference on the model and get detections\n",
    "    with torch.no_grad():\n",
    "        detections = model(input_img)\n",
    "        detections = utils.non_max_suppression(detections, 80, conf_thres, nms_thres)\n",
    "    return detections[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "ename": "RuntimeError",
     "evalue": "cuda runtime error (35) : CUDA driver version is insufficient for CUDA runtime version at /pytorch/aten/src/THC/THCGeneral.cpp:50",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-4-22c87107baf3>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[0mprev_time\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtime\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtime\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[0mimg\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mImage\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mopen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mimg_path\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 5\u001b[0;31m \u001b[0mdetections\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mdetect_image\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mimg\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      6\u001b[0m \u001b[0minference_time\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mdatetime\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtimedelta\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mseconds\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mtime\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtime\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m-\u001b[0m \u001b[0mprev_time\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      7\u001b[0m \u001b[0mprint\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0;34m'Inference Time: %s'\u001b[0m \u001b[0;34m%\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0minference_time\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;32m<ipython-input-3-e69329567a50>\u001b[0m in \u001b[0;36mdetect_image\u001b[0;34m(img)\u001b[0m\n\u001b[1;32m     12\u001b[0m     \u001b[0mimage_tensor\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mimg_transforms\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mimg\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfloat\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     13\u001b[0m     \u001b[0mimage_tensor\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mimage_tensor\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0munsqueeze_\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 14\u001b[0;31m     \u001b[0minput_img\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mVariable\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mimage_tensor\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtype\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mTensor\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     15\u001b[0m     \u001b[0;31m# run inference on the model and get detections\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     16\u001b[0m     \u001b[0;32mwith\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mno_grad\u001b[0m\u001b[0;34m(\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;31mRuntimeError\u001b[0m: cuda runtime error (35) : CUDA driver version is insufficient for CUDA runtime version at /pytorch/aten/src/THC/THCGeneral.cpp:50"
     ]
    }
   ],
   "source": [
    "# load image and get detections\n",
    "img_path = \"images/blueangels.jpg\"\n",
    "prev_time = time.time()\n",
    "img = Image.open(img_path)\n",
    "detections = detect_image(img)\n",
    "inference_time = datetime.timedelta(seconds=time.time() - prev_time)\n",
    "print ('Inference Time: %s' % (inference_time))\n",
    "\n",
    "# Get bounding-box colors\n",
    "cmap = plt.get_cmap('tab20b')\n",
    "colors = [cmap(i) for i in np.linspace(0, 1, 20)]\n",
    "\n",
    "img = np.array(img)\n",
    "plt.figure()\n",
    "fig, ax = plt.subplots(1, figsize=(12,9))\n",
    "ax.imshow(img)\n",
    "\n",
    "pad_x = max(img.shape[0] - img.shape[1], 0) * (img_size / max(img.shape))\n",
    "pad_y = max(img.shape[1] - img.shape[0], 0) * (img_size / max(img.shape))\n",
    "unpad_h = img_size - pad_y\n",
    "unpad_w = img_size - pad_x\n",
    "\n",
    "if detections is not None:\n",
    "    unique_labels = detections[:, -1].cpu().unique()\n",
    "    n_cls_preds = len(unique_labels)\n",
    "    bbox_colors = random.sample(colors, n_cls_preds)\n",
    "    # browse detections and draw bounding boxes\n",
    "    for x1, y1, x2, y2, conf, cls_conf, cls_pred in detections:\n",
    "        box_h = ((y2 - y1) / unpad_h) * img.shape[0]\n",
    "        box_w = ((x2 - x1) / unpad_w) * img.shape[1]\n",
    "        y1 = ((y1 - pad_y // 2) / unpad_h) * img.shape[0]\n",
    "        x1 = ((x1 - pad_x // 2) / unpad_w) * img.shape[1]\n",
    "        color = bbox_colors[int(np.where(unique_labels == int(cls_pred))[0])]\n",
    "        bbox = patches.Rectangle((x1, y1), box_w, box_h, linewidth=2, edgecolor=color, facecolor='none')\n",
    "        ax.add_patch(bbox)\n",
    "        plt.text(x1, y1, s=classes[int(cls_pred)], color='white', verticalalignment='top',\n",
    "                bbox={'color': color, 'pad': 0})\n",
    "plt.axis('off')\n",
    "# save image\n",
    "plt.savefig(img_path.replace(\".jpg\", \"-det.jpg\"), bbox_inches='tight', pad_inches=0.0)\n",
    "plt.show()"
   ]
  },
  {
   "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.7.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}