06_Multi_Tracking.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import Libraries\n",
    "import cv2\n",
    "import sys\n",
    "from random import randint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Tracker Types\n",
    "tracker_types = ['BOOSTING',\n",
    "                'MIL',\n",
    "                'KCF',\n",
    "                'TLD',\n",
    "                'MEADIANFLOW',\n",
    "                'GOTURN',\n",
    "                'MOSSE',\n",
    "                'CSRT']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define trackers by name\n",
    "def tracker_name(tracker_type):\n",
    "\n",
    "    # Create trackers by name with if statement\n",
    "    if tracker_type == tracker_types[0]:\n",
    "        tracker = cv2.TrackerBoosting_create()\n",
    "    elif tracker_type == tracker_types[1]:\n",
    "        tracker = cv2.TrackerMIL_create()\n",
    "    elif tracker_type == tracker_types[2]:\n",
    "        tracker = cv2.TrackerKCF_create()\n",
    "    elif tracker_type == tracker_types[3]:\n",
    "        tracker = cv2.TrackerTLD_create_()\n",
    "    elif tracker_type == tracker_types[4]:\n",
    "        tracker = cv2.TrackerMedianFlow_create()\n",
    "    elif tracker_type == tracker_types[5]:\n",
    "        tracker = cv2.TrackerGOTURN_create()\n",
    "    elif tracker_type == tracker_types[6]:\n",
    "        tracker = cv2.TrackerMOSSE_create()\n",
    "    elif tracker_type == tracker_types[7]:\n",
    "        tracker = cv2.TrackerCSRT_create()\n",
    "\n",
    "    # else statement\n",
    "    else:\n",
    "        tracker = None\n",
    "        print('No tracker found')\n",
    "        print('Choose from this trackers: ')\n",
    "        for tr in tracker_types:\n",
    "            print(tr)\n",
    "                                    \n",
    "    # return\n",
    "    return tracker                    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Default tracking algorithm MOSSE \n",
      "Available algorithms are: \n",
      "\n",
      "BOOSTING\n",
      "MIL\n",
      "KCF\n",
      "TLD\n",
      "MEADIANFLOW\n",
      "GOTURN\n",
      "MOSSE\n",
      "CSRT\n",
      "Press q to stop selecting boxes and start multitracking\n",
      "Press any key to select another box\n",
      "Press q to stop selecting boxes and start multitracking\n",
      "Press any key to select another box\n",
      "Press q to stop selecting boxes and start multitracking\n",
      "Press any key to select another box\n"
     ]
    }
   ],
   "source": [
    "if __name__ == '__main__':\n",
    "    print(\"Default tracking algorithm MOSSE \\n\" \"Available algorithms are: \\n\")\n",
    "    for ta in tracker_types:\n",
    "        print(ta)\n",
    "        \n",
    "    tracker_type = 'MOSSE'\n",
    "\n",
    "    # Create a video capture\n",
    "    cap = cv2.VideoCapture('Video/Vehicles.mp4')\n",
    "    \n",
    "    # Read first frame\n",
    "    success, frame = cap.read()\n",
    "    \n",
    "    # Quit if failure\n",
    "    if not success:\n",
    "        print('Cannot read the video')\n",
    "    \n",
    "    # Select boxes and colors\n",
    "    rects = []\n",
    "    colors = []\n",
    "\n",
    "    # While loop\n",
    "    while True:\n",
    "    \n",
    "        # draw rectangles, select ROI, open new window\n",
    "        rect_box = cv2.selectROI('MultiTrackerker', frame)\n",
    "        rects.append(rect_box)\n",
    "        colors.append((randint(64, 255), randint(64,255), randint(64,255)))\n",
    "        print('Press q to stop selecting boxes and start multitracking')\n",
    "        print('Press any key to select another box')\n",
    "        \n",
    "        \n",
    "        #close window\n",
    "        if cv2.waitKey(0) & 0xFF == 113:\n",
    "            break\n",
    "        \n",
    "    # print message\n",
    "    print(f'Selected boxes {rects}')\n",
    "\n",
    "    # Create multitracker\n",
    "    multitracker = cv2.MultiTracker_create()\n",
    "    \n",
    "    # Initialize multitracker\n",
    "    for rect_box in rects:\n",
    "        multitracker.add(tracker_name(tracker_type),\n",
    "                        frame,\n",
    "                        rect_box)\n",
    "    \n",
    "    #Video and Tracker\n",
    "    # while loop\n",
    "    while cap.isOpened():\n",
    "        success, frame = cap.read()\n",
    "        if not success:\n",
    "            break\n",
    "        \n",
    "        # update location objects\n",
    "        success, boxes = multitracker.update(frame)\n",
    "        \n",
    "        # draw the objectes tracked\n",
    "        for i, newbox in enumerate(boxes):\n",
    "            pts1 = (int(newbox[0]),\n",
    "                       int(newbox[1]))\n",
    "            pts2 = (int(newbox[0] + newbox[2]),\n",
    "                       int(newbox[1] + newbox[3]))\n",
    "            cv2.rectangle(frame, pts1, pts2, colors[i], 2, 1)\n",
    "        \n",
    "        # display frame\n",
    "        cv2.imshow('Multitracker', frame)\n",
    "    \n",
    "        # Close the frame\n",
    "        if cv2.waitKey(20) & 0xFF == 27:\n",
    "            break    \n",
    "    \n",
    "# Release and Destroy\n",
    "cap.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.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}