Vehicle_Speed_Estimation.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "## Vehicle Speed Estimation\n",
        "Estimating speed of vehicle with YOLOv5 and ByteTracker"
      ],
      "metadata": {
        "id": "o8JkIxCWeXfJ"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Installing reqruied packages & environment setup"
      ],
      "metadata": {
        "id": "9urOMuIxeeh3"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "NCaLI8Ped2F4"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "HOME = os.getcwd()"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Clones the YOLOv5 repository from GitHub using Git\n",
        "!git clone https://github.com/ultralytics/yolov5  # clone\n",
        "\n",
        "# Changes the current working directory to yolov5\n",
        "%cd yolov5\n",
        "\n",
        "# Installs required Python packages using pip\n",
        "!pip install -r requirements.txt  # install"
      ],
      "metadata": {
        "id": "5T4LTvJNeL48"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Changes the current working directory to HOME\n",
        "%cd {HOME}\n",
        "\n",
        "# Clones the ByteTrack repository from GitHub\n",
        "!git clone https://github.com/ifzhang/ByteTrack.git\n",
        "\n",
        "# Installs the required Python packages using pip3\n",
        "!cd ByteTrack && pip3 install -r requirements.txt\n",
        "\n",
        "# Runs the setup.py script using python3 to install the ByteTrack package\n",
        "!cd ByteTrack && python3 setup.py develop\n",
        "\n",
        "# Installs additional Python packages using pip\n",
        "!pip install cython_bbox loguru lap\n",
        "\n",
        "# Installs the onemetric package using pip, with the --quiet flag to suppress output\n",
        "!pip install onemetric --quiet\n"
      ],
      "metadata": {
        "id": "pKq0XgfeeL16"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Imports the 'sys' module, which provides access to some variables used or maintained by the interpreter and to functions that interact strongly with the interpreter.\n",
        "import sys\n",
        "\n",
        "# Appends a path to the system path, which is a list of directories that the interpreter will search for modules when importing them.\n",
        "# The path that is appended is constructed using an f-string that includes the value of HOME.\n",
        "sys.path.append(f\"{HOME}/ByteTrack\")"
      ],
      "metadata": {
        "id": "yAoNHxxmenFr"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Importing requirec packages\n"
      ],
      "metadata": {
        "id": "LOXTFM87euD2"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import cv2\n",
        "from dataclasses import dataclass\n",
        "from yolox.tracker.byte_tracker import BYTETracker, STrack\n",
        "from onemetric.cv.utils.iou import box_iou_batch\n",
        "import time\n",
        "import numpy as np"
      ],
      "metadata": {
        "id": "9PfMijVGenC9"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Accessories & Initializing..."
      ],
      "metadata": {
        "id": "9EDfFYdue8ju"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Define the arguments for BYTETracker\n",
        "@dataclass(frozen=True)\n",
        "class BYTETrackerArgs:\n",
        "    track_thresh: float = 0.25\n",
        "    track_buffer: int = 30\n",
        "    match_thresh: float = 0.8\n",
        "    aspect_ratio_thresh: float = 3.0\n",
        "    min_box_area: float = 1.0\n",
        "    mot20: bool = False\n",
        "\n",
        "# Define a function to calculate the speed of a vehicle\n",
        "def calculate_speed(previous_centroid, current_centroid, previous_frame_time, current_frame_time):\n",
        "    # Calculate the distance traveled by the vehicle\n",
        "    distance = np.linalg.norm(current_centroid - previous_centroid)\n",
        "\n",
        "    # Calculate the time taken by the vehicle to travel the distance\n",
        "    time_delta = current_frame_time - previous_frame_time\n",
        "\n",
        "    # Calculate the speed of the vehicle in pixels per second\n",
        "    speed = distance / time_delta\n",
        "\n",
        "    # Return the speed in kilometers per hour\n",
        "    return speed * 3.6\n",
        "  \n",
        "#Initialize previous frame time \n",
        "previous_frame_time = None\n",
        "# Initiate the BYTETracker object\n",
        "tracker = BYTETracker(BYTETrackerArgs)\n",
        "\n",
        "# Define the vehicle classes to be tracked\n",
        "vehicle_classes = ['car', 'bus', 'truck']\n",
        "\n",
        "# Get the names of the classes from the YOLOv5 model\n",
        "vehicle_names = model.names\n",
        "\n",
        "# Initialize a dictionary to store the bounding boxes for each track ID in the previous frame\n",
        "prevoius_online_tracker  = {}"
      ],
      "metadata": {
        "id": "wi7iQPhkfCEC"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Loading model"
      ],
      "metadata": {
        "id": "NVS6akwae2Mt"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Load the YOLOv5 model\n",
        "model = torch.hub.load(\"ultralytics/yolov5\", \"yolov5s\")  # or yolov5n - yolov5x6, custom\n",
        "\n",
        "nmodel = torch.hub.load(\"ultralytics/yolov5\", 'custom', \"/Users/auk/Downloads/DataProjects/speed_estimation/best.pt\")  # or yolov5n - yolov5x6, custom"
      ],
      "metadata": {
        "id": "D_YrCYsnem_7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Input data"
      ],
      "metadata": {
        "id": "TSX0uRVpfPIK"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Open the video file for processing\n",
        "cap = cv2.VideoCapture('/Users/auk/Downloads/DataProjects/speed_estimation/b.mp4')"
      ],
      "metadata": {
        "id": "lZ_VNHPDfUYo"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Input Processing & Speed estimation"
      ],
      "metadata": {
        "id": "O8Ffb9tgfWbG"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Loop through each frame of the video\n",
        "while True:\n",
        "    # Read the frame\n",
        "    ret, frame = cap.read()\n",
        "\n",
        "    if not ret:\n",
        "        break\n",
        "\n",
        "    # Perform object detection with YOLOv5\n",
        "    results = model(frame)\n",
        "\n",
        "    # Update the tracker with the detection results for the vehicle classes\n",
        "    for det in results.pred[0]:\n",
        "        class_idx = int(det[-1])\n",
        "        conf = float(det[-2])\n",
        "        if vehicle_names[class_idx] in vehicle_classes and conf > 0.7:\n",
        "            online_tracker= tracker.update(results.pred[0], img_info=frame.shape, img_size=frame.shape)\n",
        "\n",
        "    # Draw the bounding boxes and track IDs on the frame\n",
        "    for obj in online_tracker:\n",
        "        x1, y1, x2, y2 = map(int, obj.tlbr)\n",
        "        current_frame_time = time.time()\n",
        "\n",
        "        if previous_frame_time == None:\n",
        "            break\n",
        "        if obj.track_id in prevoius_online_tracker.keys():\n",
        "            px1, py1, px2, py2 = prevoius_online_tracker[obj.track_id]\n",
        "            previous_centroid = np.array([(px1+px2)/2, (py1+py2)/2])\n",
        "            current_centroid = np.array([(x1+x2)/2, (y1+y2)/2])\n",
        "            speed = calculate_speed(previous_centroid, current_centroid, previous_frame_time, current_frame_time)\n",
        "            if speed > 0:\n",
        "                cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 1)\n",
        "                cv2.putText(frame, f'Speed: {speed:.2f} km/h', (x1, y1-5), cv2.FONT_HERSHEY_SIMPLEX, 0.3, (0, 255, 0), 1)\n",
        "\n",
        "        prevoius_online_tracker[obj.track_id] = [x1, y1, x2, y2]\n",
        "        \n",
        "\n",
        "    previous_frame_time = current_frame_time\n",
        "\n",
        "    # Display the processed frame\n",
        "    cv2.imshow('Speed Estimator', frame)\n",
        "\n",
        "    # Exit the loop if the 's' key is pressed\n",
        "    if cv2.waitKey(1) == ord('s'):\n",
        "        break\n",
        "\n",
        "# Release the resources\n",
        "cap.release()\n",
        "cv2.destroyAllWindows()\n"
      ],
      "metadata": {
        "id": "e5gtWysxeLzM"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}