VehicleCounterYOLO.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import imutils\n",
    "import time\n",
    "import cv2\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] loading YOLO from disk...\n",
      "[INFO] 200 total frames in video\n",
      "Total vehicles in image 8\n",
      "Each vehicles count in image {'car': 8}\n",
      "[INFO] single frame took 0.9016 seconds\n",
      "[INFO] estimated total time to finish: 180.3188\n",
      "Total vehicles in image 8\n",
      "Each vehicles count in image {'car': 8}\n",
      "Total vehicles in image 8\n",
      "Each vehicles count in image {'car': 8}\n",
      "Total vehicles in image 6\n",
      "Each vehicles count in image {'car': 6}\n",
      "Total vehicles in image 8\n",
      "Each vehicles count in image {'car': 8}\n",
      "Total vehicles in image 8\n",
      "Each vehicles count in image {'car': 8}\n",
      "Total vehicles in image 12\n",
      "Each vehicles count in image {'car': 12}\n",
      "Total vehicles in image 14\n",
      "Each vehicles count in image {'car': 14}\n",
      "Total vehicles in image 14\n",
      "Each vehicles count in image {'car': 14}\n",
      "Total vehicles in image 13\n",
      "Each vehicles count in image {'car': 13}\n",
      "Total vehicles in image 11\n",
      "Each vehicles count in image {'car': 11}\n",
      "Total vehicles in image 14\n",
      "Each vehicles count in image {'car': 14}\n",
      "Total vehicles in image 14\n",
      "Each vehicles count in image {'car': 14}\n",
      "Total vehicles in image 11\n",
      "Each vehicles count in image {'car': 11}\n",
      "Total vehicles in image 13\n",
      "Each vehicles count in image {'car': 13}\n",
      "Total vehicles in image 12\n",
      "Each vehicles count in image {'car': 12}\n",
      "Total vehicles in image 12\n",
      "Each vehicles count in image {'car': 12}\n",
      "Total vehicles in image 12\n",
      "Each vehicles count in image {'car': 12}\n",
      "Total vehicles in image 11\n",
      "Each vehicles count in image {'car': 11}\n",
      "Total vehicles in image 12\n",
      "Each vehicles count in image {'car': 12}\n",
      "Total vehicles in image 9\n",
      "Each vehicles count in image {'car': 9}\n",
      "Total vehicles in image 11\n",
      "Each vehicles count in image {'car': 11}\n",
      "Total vehicles in image 9\n",
      "Each vehicles count in image {'car': 9}\n",
      "Total vehicles in image 12\n",
      "Each vehicles count in image {'car': 12}\n",
      "Total vehicles in image 11\n",
      "Each vehicles count in image {'car': 11}\n",
      "Total vehicles in image 11\n",
      "Each vehicles count in image {'car': 11}\n",
      "Total vehicles in image 8\n",
      "Each vehicles count in image {'car': 8}\n",
      "Total vehicles in image 11\n",
      "Each vehicles count in image {'car': 11}\n",
      "Total vehicles in image 10\n",
      "Each vehicles count in image {'car': 10}\n",
      "Total vehicles in image 11\n",
      "Each vehicles count in image {'car': 11}\n",
      "Total vehicles in image 9\n",
      "Each vehicles count in image {'car': 9}\n",
      "Total vehicles in image 11\n",
      "Each vehicles count in image {'car': 11}\n",
      "Total vehicles in image 11\n",
      "Each vehicles count in image {'car': 11}\n",
      "Total vehicles in image 10\n",
      "Each vehicles count in image {'car': 10}\n",
      "Total vehicles in image 11\n",
      "Each vehicles count in image {'car': 11}\n",
      "Total vehicles in image 10\n",
      "Each vehicles count in image {'car': 10}\n",
      "Total vehicles in image 10\n",
      "Each vehicles count in image {'car': 10}\n",
      "Total vehicles in image 11\n",
      "Each vehicles count in image {'car': 11}\n",
      "Total vehicles in image 9\n",
      "Each vehicles count in image {'car': 9}\n",
      "Total vehicles in image 10\n",
      "Each vehicles count in image {'car': 10}\n",
      "Total vehicles in image 8\n",
      "Each vehicles count in image {'car': 8}\n",
      "Total vehicles in image 10\n",
      "Each vehicles count in image {'car': 10}\n",
      "Total vehicles in image 11\n",
      "Each vehicles count in image {'car': 11}\n",
      "Total vehicles in image 14\n",
      "Each vehicles count in image {'car': 14}\n",
      "Total vehicles in image 12\n",
      "Each vehicles count in image {'car': 12}\n",
      "Total vehicles in image 11\n",
      "Each vehicles count in image {'car': 11}\n",
      "Total vehicles in image 9\n",
      "Each vehicles count in image {'car': 9}\n",
      "Total vehicles in image 8\n",
      "Each vehicles count in image {'car': 8}\n",
      "Total vehicles in image 9\n",
      "Each vehicles count in image {'car': 9}\n",
      "Total vehicles in image 11\n",
      "Each vehicles count in image {'car': 11}\n",
      "Total vehicles in image 12\n",
      "Each vehicles count in image {'car': 12}\n",
      "Total vehicles in image 11\n",
      "Each vehicles count in image {'car': 11}\n",
      "Total vehicles in image 12\n",
      "Each vehicles count in image {'car': 12}\n",
      "Total vehicles in image 11\n",
      "Each vehicles count in image {'car': 11}\n",
      "Total vehicles in image 12\n",
      "Each vehicles count in image {'car': 12}\n",
      "Total vehicles in image 14\n",
      "Each vehicles count in image {'car': 14}\n",
      "Total vehicles in image 15\n",
      "Each vehicles count in image {'car': 15}\n",
      "Total vehicles in image 12\n",
      "Each vehicles count in image {'car': 12}\n",
      "Total vehicles in image 11\n",
      "Each vehicles count in image {'car': 11}\n",
      "[INFO] cleaning up...\n"
     ]
    }
   ],
   "source": [
    "conf = 0.5 #minimum probability to filter weak detections\n",
    "thd = 0.3 #threshold non-maxima suppression\n",
    "\n",
    "classname = []\n",
    "list_of_vehicles = [\"bicycle\",\"car\",\"motorbike\",\"bus\",\"truck\"]\n",
    "def get_vehicle_count(boxes, class_names):\n",
    "    total_vehicle_count = 0 # total vechiles present in the image\n",
    "    dict_vehicle_count = {} # dictionary with count of each distinct vehicles detected\n",
    "    for i in range(len(boxes)):\n",
    "        class_name = class_names[i]\n",
    "        # print(i,\".\",class_name)\n",
    "        if(class_name in list_of_vehicles):\n",
    "            total_vehicle_count += 1\n",
    "            dict_vehicle_count[class_name] = dict_vehicle_count.get(class_name,0) + 1\n",
    "\n",
    "    return total_vehicle_count, dict_vehicle_count\n",
    "\n",
    "# load the COCO class labels our YOLO model was trained on\n",
    "labelsPath = os.path.sep.join([\"./yolo-coco\", \"coco.names\"])\n",
    "LABELS = open(labelsPath).read().strip().split(\"\\n\")\n",
    "\n",
    "# initialize a list of colors to represent each possible class label\n",
    "np.random.seed(42)\n",
    "COLORS = np.random.randint(0, 255, size=(len(LABELS), 3),dtype=\"uint8\")\n",
    "\n",
    "\n",
    "# derive the paths to the YOLO weights and model configuration\n",
    "weightsPath = os.path.sep.join([\"./yolo-coco\", \"yolov3.weights\"])\n",
    "configPath = os.path.sep.join([\"./yolo-coco\", \"yolov3.cfg\"])\n",
    "\n",
    "# load our YOLO object detector trained on COCO dataset (80 classes)\n",
    "# and determine only the *output* layer names that we need from YOLO\n",
    "print(\"[INFO] loading YOLO from disk...\")\n",
    "net = cv2.dnn.readNetFromDarknet(configPath, weightsPath)\n",
    "ln = net.getLayerNames()\n",
    "ln = [ln[i[0] - 1] for i in net.getUnconnectedOutLayers()]\n",
    "\n",
    "# initialize the video stream, pointer to output video file, and\n",
    "# frame dimensions\n",
    "vs = cv2.VideoCapture(\"./Video_Data_Samples/video1.avi\")\n",
    "writer = None\n",
    "(W, H) = (None, None)\n",
    "\n",
    "# try to determine the total number of frames in the video file\n",
    "try:\n",
    "    prop = cv2.cv.CV_CAP_PROP_FRAME_COUNT if imutils.is_cv2() \\\n",
    "        else cv2.CAP_PROP_FRAME_COUNT\n",
    "    total = 200\n",
    "    print(\"[INFO] {} total frames in video\".format(total))\n",
    "\n",
    "# an error occurred while trying to determine the total\n",
    "# number of frames in the video file\n",
    "except:\n",
    "    print(\"[INFO] could not determine # of frames in video\")\n",
    "    print(\"[INFO] no approx. completion time can be provided\")\n",
    "    total = -1\n",
    "\n",
    "# loop over frames from the video file stream\n",
    "\n",
    "list_of_vehicles = [\"car\",\"bus\",\"motorbike\",\"truck\",\"bicycle\"]\n",
    "\n",
    "while True:\n",
    "    # read the next frame from the file\n",
    "    (grabbed, frame) = vs.read()\n",
    "\n",
    "    # if the frame was not grabbed, then we have reached the end\n",
    "    # of the stream\n",
    "    if not grabbed:\n",
    "        break\n",
    "\n",
    "    # if the frame dimensions are empty, grab them\n",
    "    if W is None or H is None:\n",
    "        (H, W) = frame.shape[:2]\n",
    "\n",
    "    # construct a blob from the input frame and then perform a forward\n",
    "    # pass of the YOLO object detector, giving us our bounding boxes\n",
    "    # and associated probabilities\n",
    "    blob = cv2.dnn.blobFromImage(frame, 1 / 255.0, (416, 416),swapRB=True, crop=False)\n",
    "    net.setInput(blob)\n",
    "    start = time.time()\n",
    "    layerOutputs = net.forward(ln)\n",
    "    end = time.time()\n",
    "\n",
    "    # initialize our lists of detected bounding boxes, confidences,\n",
    "    # and class IDs, respectively\n",
    "    boxes = []\n",
    "    confidences = []\n",
    "    classIDs = []\n",
    "\n",
    "    # loop over each of the layer outputs\n",
    "    for output in layerOutputs:\n",
    "        # loop over each of the detections\n",
    "        for detection in output:\n",
    "            # extract the class ID and confidence (i.e., probability)\n",
    "            # of the current object detection\n",
    "            scores = detection[5:]\n",
    "            classID = np.argmax(scores)\n",
    "            confidence = scores[classID]\n",
    "\n",
    "            # filter out weak predictions by ensuring the detected\n",
    "            # probability is greater than the minimum probability\n",
    "            if confidence > conf:\n",
    "                # scale the bounding box coordinates back relative to\n",
    "                # the size of the image, keeping in mind that YOLO\n",
    "                # actually returns the center (x, y)-coordinates of\n",
    "                # the bounding box followed by the boxes' width and\n",
    "                # height\n",
    "                box = detection[0:4] * np.array([W, H, W, H])\n",
    "                (centerX, centerY, width, height) = box.astype(\"int\")\n",
    "\n",
    "                # use the center (x, y)-coordinates to derive the top\n",
    "                # and and left corner of the bounding box\n",
    "                x = int(centerX - (width / 2))\n",
    "                y = int(centerY - (height / 2))\n",
    "\n",
    "                # update our list of bounding box coordinates,\n",
    "                # confidences, and class IDs\n",
    "                boxes.append([x, y, int(width), int(height)])\n",
    "                confidences.append(float(confidence))\n",
    "                classIDs.append(classID)\n",
    "                classname.append(LABELS[classID])\n",
    "\n",
    "    # apply non-maxima suppression to suppress weak, overlapping\n",
    "    # bounding boxes\n",
    "    idxs = cv2.dnn.NMSBoxes(boxes, confidences,conf,thd)\n",
    "\n",
    "\n",
    "    # ensure at least one detection exists\n",
    "    if len(idxs) > 0:\n",
    "        # loop over the indexes we are keeping\n",
    "        for i in idxs.flatten():\n",
    "            # extract the bounding box coordinates\n",
    "            (x, y) = (boxes[i][0], boxes[i][1])\n",
    "            (w, h) = (boxes[i][2], boxes[i][3])\n",
    "\n",
    "            # draw a bounding box rectangle and label on the frame\n",
    "            color = [int(c) for c in COLORS[classIDs[i]]]\n",
    "            cv2.rectangle(frame, (x, y), (x + w, y + h), color, 2)\n",
    "            text = \"{}: {:.4f}\".format(LABELS[classIDs[i]],confidences[i])\n",
    "            cv2.putText(frame, text, (x, y - 5),cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)    \n",
    "            \n",
    "\n",
    "    cv2.imshow('Frame', cv2.resize(frame, (400, 400))) \n",
    "    # check if the video writer is None\n",
    "    total_vehicles, each_vehicle = get_vehicle_count(boxes, classname)\n",
    "    print(\"Total vehicles in image\", total_vehicles)\n",
    "    print(\"Each vehicles count in image\", each_vehicle)\n",
    "    if writer is None:\n",
    "        # initialize our video writer\n",
    "        fourcc = cv2.VideoWriter_fourcc(*\"MJPG\")\n",
    "        writer = cv2.VideoWriter(\"./output.avi\", fourcc, 30,(frame.shape[1], frame.shape[0]), True)\n",
    "\n",
    "\n",
    "        # some information on processing single frame\n",
    "        if total > 0:\n",
    "            elap = (end - start)\n",
    "            print(\"[INFO] single frame took {:.4f} seconds\".format(elap))\n",
    "            print(\"[INFO] estimated total time to finish: {:.4f}\".format(elap * total))\n",
    "\n",
    "    # write the output frame to disk\n",
    "    writer.write(frame)\n",
    "    if cv2.waitKey(1) & 0xFF == ord('q'):\n",
    "        break    \n",
    "# release the file pointers\n",
    "writer.release()\n",
    "vs.release()\n",
    "cv2.destroyAllWindows()"
   ]
  },
  {
   "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.6.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}