example-scripts/tracking/IoU/IoU.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "c6175afe",
   "metadata": {},
   "source": [
    "## IOU Tracking"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c15e649d",
   "metadata": {},
   "outputs": [],
   "source": [
    "#!/usr/bin/env python\n",
    "# -*-coding:utf-8 -*-\n",
    "\"\"\"\n",
    "  ████\n",
    "██    ██   Datature\n",
    "  ██  ██   Powering Breakthrough AI\n",
    "    ██\n",
    " \n",
    "@File    :   IoU.ipynb\n",
    "@Author  :   Keechin Goh\n",
    "@Version :   1.0\n",
    "@Contact :   hello@datature.io\n",
    "@License :   Apache License 2.0\n",
    "@Desc    :   IoU Tracking Script\n",
    "\"\"\""
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "2c62ea26",
   "metadata": {},
   "source": [
    "### Introduction\n",
    "This jupyter notebook runs a Norfair Algorithm atop of various object detection (bounding box) models to track objects in video files."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bcd91dc5",
   "metadata": {},
   "source": [
    "## Install Python Dependencies\n",
    "\n",
    "Python Version Used: 3.7 =< version =< 3.9"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "afe4a404",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -U tensorflow==2.5.0\n",
    "!pip install -U numpy==1.18.5\n",
    "!pip install -U Pillow==7.2.0\n",
    "!pip install -U opencv-python==4.5.1.48\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d6d81c47",
   "metadata": {},
   "source": [
    "## Import Python Libraries\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c68d01ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import cv2\n",
    "import time\n",
    "import os.path\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "\n",
    "from PIL import Image\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6388c29a",
   "metadata": {},
   "source": [
    "## Load Label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1a998d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_label(model_path):\n",
    "    \"\"\"Reads label map in the format of .pbtxt and parse into dictionary\n",
    "\n",
    "    Args:\n",
    "        label_map_path: the file path to the label_map\n",
    "\n",
    "    Returns:\n",
    "        dictionary with the format of {label_index: {'id': label_index, 'name': label_name}}\n",
    "    \"\"\"\n",
    "    label_map_path = os.path.join(model_path, \"label_map.pbtxt\")\n",
    "\n",
    "    if os.path.exists(label_map_path) is False:\n",
    "        raise FileNotFoundError(\"No valid label map found.\")\n",
    "\n",
    "    label_map = {}\n",
    "\n",
    "    with open(label_map_path, \"r\") as label_file:\n",
    "        for line in label_file:\n",
    "            if \"id\" in line:\n",
    "                label_index = int(line.split(\":\")[-1])\n",
    "                label_name = next(label_file).split(\":\")[-1].strip().strip(\"'\")\n",
    "                label_map[label_index] = {\"id\": label_index, \"name\": label_name}\n",
    "\n",
    "    return label_map\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee960ecc",
   "metadata": {},
   "source": [
    "## Load Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc7c12f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_image_into_numpy_array(img, height, width):\n",
    "    \"\"\"Load an image from base64 into a numpy array.\n",
    "\n",
    "    Puts image into numpy array to feed into tensorflow graph.\n",
    "    Note that by convention we put it into a numpy array with shape\n",
    "    (height, width, channels), where channels=3 for RGB.\n",
    "\n",
    "    Args:\n",
    "        image_base64: image in base64 format\n",
    "        height: height of image\n",
    "        width: width of image\n",
    "\n",
    "    Returns:\n",
    "        uint8 numpy array with shape (img_height, img_width, 3)\n",
    "    \"\"\"\n",
    "    image_shape = np.asarray(img).shape\n",
    "\n",
    "    image_resized = img.resize((height, width))\n",
    "    return np.array(image_resized), (image_shape[0], image_shape[1])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "315a706a",
   "metadata": {},
   "source": [
    "## Prediction Functions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a0a37e4",
   "metadata": {},
   "source": [
    "### Bounding Box Tensorflow Prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8d0993f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def bbox_tf_predition(img, trained_model, size, threshold):\n",
    "    \"\"\"Run prediction on image in base64 format with trained tf bounding box model and label.\n",
    "\n",
    "    Args:\n",
    "        img: input image\n",
    "        trained_model: loaded tensorflow model.\n",
    "        size: size of image to load.\n",
    "        threshold: confidence threshold.\n",
    "\n",
    "    Returns:\n",
    "        list for detection in the format of\n",
    "        [(boxes,classes,scores)...]\n",
    "        boxes: Array of array contains normalized boxes coorindates,\n",
    "        classes: Array of numbers representing class,\n",
    "        scores: Array of numbers representing confidence scores\n",
    "    \"\"\"\n",
    "\n",
    "    ## Run prediction\n",
    "    height, width = size.split(\",\")\n",
    "\n",
    "    ## Returned original_shape is in the format of width, height\n",
    "    image_resized, origi_shape = load_image_into_numpy_array(\n",
    "        img, int(height), int(width)\n",
    "    )\n",
    "\n",
    "    ## The input needs to be a tensor, convert it using `tf.convert_to_tensor`.\n",
    "    input_tensor = tf.convert_to_tensor(image_resized)\n",
    "\n",
    "    ## The model expects a batch of images, so add an axis with `tf.newaxis`.\n",
    "    input_tensor = input_tensor[tf.newaxis, ...]\n",
    "\n",
    "    ## Feed image into model\n",
    "    detections_output = trained_model(input_tensor)\n",
    "\n",
    "    num_detections = int(detections_output.pop(\"num_detections\"))\n",
    "    detections = {\n",
    "        key: value[0, :num_detections].numpy()\n",
    "        for key, value in detections_output.items()\n",
    "    }\n",
    "    detections[\"num_detections\"] = num_detections\n",
    "\n",
    "    ## Filter out predictions below threshold\n",
    "    indexes = np.where(detections[\"detection_scores\"] > float(threshold))\n",
    "\n",
    "    ## Extract predictions\n",
    "    bboxes = detections[\"detection_boxes\"][indexes]\n",
    "    classes = detections[\"detection_classes\"][indexes].astype(np.int64)\n",
    "    scores = detections[\"detection_scores\"][indexes]\n",
    "    detection_info = {\n",
    "        \"boxes\": bboxes.tolist(),\n",
    "        \"classes\": classes.tolist(),\n",
    "        \"scores\": scores.tolist(),\n",
    "    }\n",
    "    if len(detection_info[\"boxes\"]) == 0:\n",
    "        return [{\"bbox\": [], \"score\": -1, \"class\": -1}]\n",
    "    else:\n",
    "        detection_out = []\n",
    "        for i in range(0, len(detection_info[\"boxes\"])):\n",
    "            temp_box = [0, 0, 0, 0]\n",
    "            temp_box[0] = detection_info[\"boxes\"][i][1] * origi_shape[1]\n",
    "            temp_box[1] = detection_info[\"boxes\"][i][0] * origi_shape[0]\n",
    "            temp_box[2] = detection_info[\"boxes\"][i][3] * origi_shape[1]\n",
    "            temp_box[3] = detection_info[\"boxes\"][i][2] * origi_shape[0]\n",
    "\n",
    "            detection_out.append(\n",
    "                {\n",
    "                    \"bbox\": temp_box,\n",
    "                    \"score\": detection_info[\"scores\"][i],\n",
    "                    \"class\": detection_info[\"classes\"][i],\n",
    "                }\n",
    "            )\n",
    "\n",
    "    return detection_out\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a57ae95",
   "metadata": {},
   "source": [
    "## IOU Tracking"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08a70e14",
   "metadata": {},
   "outputs": [],
   "source": [
    "def iou(bbox1, bbox2):\n",
    "    \"\"\"\n",
    "    Calculates the intersection-over-union of two bounding boxes.\n",
    "    Args:\n",
    "        bbox1 (numpy.array, list of floats): bounding box in format x1,y1,x2,y2.\n",
    "        bbox2 (numpy.array, list of floats): bounding box in format x1,y1,x2,y2.\n",
    "    Returns:\n",
    "        int: intersection-over-onion of bbox1, bbox2\n",
    "    \"\"\"\n",
    "\n",
    "    bbox1 = [float(x) for x in bbox1]\n",
    "    bbox2 = [float(x) for x in bbox2]\n",
    "\n",
    "    (x0_1, y0_1, x1_1, y1_1) = bbox1\n",
    "    (x0_2, y0_2, x1_2, y1_2) = bbox2\n",
    "\n",
    "    # get the overlap rectangle\n",
    "    overlap_x0 = max(x0_1, x0_2)\n",
    "    overlap_y0 = max(y0_1, y0_2)\n",
    "    overlap_x1 = min(x1_1, x1_2)\n",
    "    overlap_y1 = min(y1_1, y1_2)\n",
    "\n",
    "    # check if there is an overlap\n",
    "    if overlap_x1 - overlap_x0 <= 0 or overlap_y1 - overlap_y0 <= 0:\n",
    "        return 0\n",
    "\n",
    "    # if yes, calculate the ratio of the overlap to each ROI size and the unified size\n",
    "    size_1 = (x1_1 - x0_1) * (y1_1 - y0_1)\n",
    "    size_2 = (x1_2 - x0_2) * (y1_2 - y0_2)\n",
    "    size_intersection = (overlap_x1 - overlap_x0) * (overlap_y1 - overlap_y0)\n",
    "    size_union = size_1 + size_2 - size_intersection\n",
    "\n",
    "    return size_intersection / size_union\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53ae0a45",
   "metadata": {},
   "outputs": [],
   "source": [
    "def track_iou(detections, sigma_l, sigma_h, sigma_iou, t_min):\n",
    "    \"\"\"\n",
    "    Simple IOU based tracker.\n",
    "    Args:\n",
    "         detections (list): list of detections per frame\n",
    "         sigma_l (float): low detection threshold.\n",
    "         sigma_h (float): high detection threshold.\n",
    "         sigma_iou (float): IOU threshold.\n",
    "         t_min (float): minimum track length in frames.\n",
    "    Returns:\n",
    "        list: list of tracks.\n",
    "    \"\"\"\n",
    "\n",
    "    tracks_active = []\n",
    "    tracks_finished = []\n",
    "    track_id = 1\n",
    "    for frame_num, detections_frame in enumerate(detections, start=1):\n",
    "        # apply low threshold to detections\n",
    "        dets = [det for det in detections_frame if det[\"score\"] >= sigma_l]\n",
    "        updated_tracks = []\n",
    "        for track in tracks_active:\n",
    "\n",
    "            if len(dets) > 0:\n",
    "                # get det with highest iou\n",
    "                best_match = max(\n",
    "                    dets, key=lambda x: iou(track[\"bboxes\"][-1], x[\"bbox\"])\n",
    "                )\n",
    "                if iou(track[\"bboxes\"][-1], best_match[\"bbox\"]) >= sigma_iou:\n",
    "                    track[\"bboxes\"].append(best_match[\"bbox\"])\n",
    "                    track[\"max_score\"] = max(track[\"max_score\"], best_match[\"score\"])\n",
    "                    track[\"frame_num\"].append(frame_num)\n",
    "\n",
    "                    updated_tracks.append(track)\n",
    "\n",
    "                    # remove from best matching detection from detections\n",
    "                    del dets[dets.index(best_match)]\n",
    "\n",
    "            # if track was not updated\n",
    "            if len(updated_tracks) == 0 or track is not updated_tracks[-1]:\n",
    "                # finish track when the conditions are met\n",
    "                if track[\"max_score\"] >= sigma_h and len(track[\"bboxes\"]) >= t_min:\n",
    "                    track[\"frame_num\"].append(frame_num)\n",
    "                    track[\"track_id\"] = track_id\n",
    "                    track_id += 1\n",
    "                    tracks_finished.append(track)\n",
    "\n",
    "        # create new tracks\n",
    "        new_tracks = [\n",
    "            {\n",
    "                \"bboxes\": [det[\"bbox\"]],\n",
    "                \"max_score\": det[\"score\"],\n",
    "                \"start_frame\": frame_num,\n",
    "                \"frame_num\": [],\n",
    "                \"classes\": det[\"class\"],\n",
    "            }\n",
    "            for det in dets\n",
    "        ]\n",
    "        tracks_active = updated_tracks + new_tracks\n",
    "\n",
    "    # finish all remaining active tracks\n",
    "    for track in tracks_active:\n",
    "        if track[\"max_score\"] >= sigma_h and len(track[\"bboxes\"]) >= t_min:\n",
    "            # track[\"frame_num\"].append(frame_num)\n",
    "            track[\"track_id\"] = track_id\n",
    "            track_id += 1\n",
    "            tracks_finished.append(track)\n",
    "\n",
    "    return tracks_finished, frame_num\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8dfe0704",
   "metadata": {},
   "outputs": [],
   "source": [
    "def IOU_track(\n",
    "    video_path, model_path, size, threshold, output_format, output_vid_trk_path, COLORS\n",
    "):\n",
    "\n",
    "    # load model\n",
    "    trained_model = tf.saved_model.load(os.path.join(model_path, \"saved_model\"))\n",
    "    # load map\n",
    "    category_index = load_label(model_path)\n",
    "\n",
    "    # read in video\n",
    "    vid = cv2.VideoCapture(video_path)\n",
    "\n",
    "    # initialize video writer\n",
    "    width = int(vid.get(cv2.CAP_PROP_FRAME_WIDTH))\n",
    "    height = int(vid.get(cv2.CAP_PROP_FRAME_HEIGHT))\n",
    "    fps = int(vid.get(cv2.CAP_PROP_FPS))\n",
    "    codec = cv2.VideoWriter_fourcc(*output_format)\n",
    "    out = cv2.VideoWriter(output_vid_trk_path, codec, fps, (width, height))\n",
    "\n",
    "    # initail Tracker\n",
    "    sigma_l = 0\n",
    "    sigma_h = float(threshold)\n",
    "    sigma_iou = 0.5\n",
    "    t_min = 2\n",
    "\n",
    "    # Generate detection list for each frame\n",
    "    detection = []\n",
    "    while True:\n",
    "        # Read Video\n",
    "        (grabbed, frame) = vid.read()\n",
    "        if not grabbed:\n",
    "            break\n",
    "\n",
    "        # Do detection on every video frame\n",
    "        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)\n",
    "        prediction = bbox_tf_predition(\n",
    "            Image.fromarray(frame), trained_model, size, threshold\n",
    "        )\n",
    "        detection.append(prediction)\n",
    "    vid.release()\n",
    "    print(\"Prediction completed!\")\n",
    "    tracks_ori = track_iou(detection, sigma_l, sigma_h, sigma_iou, t_min)\n",
    "    tracks = []\n",
    "    for i in range(tracks_ori[1]):\n",
    "        tracks.append([])\n",
    "    for i in tracks_ori[0]:\n",
    "        for j in range(len(i[\"frame_num\"])):\n",
    "            tracks[i[\"frame_num\"][j] - 1].append(\n",
    "                {\n",
    "                    \"box\": i[\"bboxes\"][j],\n",
    "                    \"track_id\": i[\"track_id\"],\n",
    "                    \"class\": i[\"classes\"],\n",
    "                }\n",
    "            )\n",
    "\n",
    "    idx = 0\n",
    "    # read in video\n",
    "    vid = cv2.VideoCapture(video_path)\n",
    "    while True:\n",
    "        # Read Video\n",
    "        (grabbed, frame) = vid.read()\n",
    "        if not grabbed:\n",
    "            break\n",
    "\n",
    "        # Do detection on every video frame\n",
    "        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)\n",
    "        if tracks[idx] == []:\n",
    "            result = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)\n",
    "            out.write(result)\n",
    "            print(\"Frame \" + str(idx) + \" is written!\")\n",
    "            idx += 1\n",
    "            continue\n",
    "\n",
    "        # draw tracks\n",
    "        for track in tracks[idx]:\n",
    "            bbox = track[\"box\"]\n",
    "            class_name = category_index[int(track[\"class\"])][\"name\"].strip('\"')\n",
    "            color = tuple(int(c) for c in COLORS[int(track[\"track_id\"]) % len(COLORS)])\n",
    "\n",
    "            # draw bbox on screen\n",
    "            cv2.rectangle(\n",
    "                frame,\n",
    "                (\n",
    "                    int(bbox[0]),\n",
    "                    int(bbox[1]),\n",
    "                ),\n",
    "                (\n",
    "                    int(bbox[2]),\n",
    "                    int(bbox[3]),\n",
    "                ),\n",
    "                color,\n",
    "                int((width + height) / 600),\n",
    "            )\n",
    "            ## Draw label background\n",
    "            cv2.rectangle(\n",
    "                frame,\n",
    "                (\n",
    "                    int(bbox[0]),\n",
    "                    int(bbox[3]),\n",
    "                ),\n",
    "                (\n",
    "                    int(bbox[2]),\n",
    "                    int(bbox[3] + int((width + height) / 108)),\n",
    "                ),\n",
    "                color,\n",
    "                -1,\n",
    "            )\n",
    "            cv2.putText(\n",
    "                frame,\n",
    "                class_name,\n",
    "                (int(bbox[0]), int(bbox[3] + int((width + height) / 300))),\n",
    "                cv2.FONT_HERSHEY_SIMPLEX,\n",
    "                (width + height) / 7500,\n",
    "                (0, 0, 0),\n",
    "                int((width + height) / 3000),\n",
    "                cv2.LINE_AA,\n",
    "            )\n",
    "            cv2.putText(\n",
    "                frame,\n",
    "                \"ID: \" + str(track[\"track_id\"]),\n",
    "                (int(bbox[0]), int(bbox[3] + int((width + height) / 136))),\n",
    "                cv2.FONT_HERSHEY_SIMPLEX,\n",
    "                (width + height) / 7500,\n",
    "                (0, 0, 0),\n",
    "                int((width + height) / 3000),\n",
    "                cv2.LINE_AA,\n",
    "            )\n",
    "\n",
    "        result = np.asarray(frame)\n",
    "        result = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)\n",
    "        # save video\n",
    "        out.write(result)\n",
    "        print(\"Frame \" + str(idx) + \" is written!\")\n",
    "        idx += 1\n",
    "\n",
    "    vid.release()\n",
    "    out.release()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4fcd3ef5",
   "metadata": {},
   "source": [
    "## Run IOU for Video"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "616de082",
   "metadata": {},
   "outputs": [],
   "source": [
    "video_path = \"./pexels-rodnae-productions-5699922.mp4\"\n",
    "model_path = \"./model\"\n",
    "output_vid_trk_path = \"./output_\" + video_path.split(\"./\")[1].split(\".\")[0] + \".mp4\"\n",
    "size = \"320,320\"\n",
    "threshold = \"0.7\"\n",
    "output_format = \"mp4v\"\n",
    "COLORS = np.random.randint(0, 255, size=(200, 3))\n",
    "\n",
    "start = time.time()\n",
    "output_vid = IOU_track(\n",
    "    video_path, model_path, size, threshold, output_format, output_vid_trk_path, COLORS\n",
    ")\n",
    "end = time.time()\n",
    "print(\"Tracking by IOU for \" + video_path + \" takes \" + str(int(end - start)) + \"s.\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ab6614c",
   "metadata": {},
   "source": [
    "## Display Video"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0f77795",
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import display, HTML\n",
    "from base64 import b64encode\n",
    "\n",
    "\n",
    "def display_video(path):\n",
    "    mp4 = open(path, \"rb\").read()\n",
    "    data_url = \"data:video/mp4;base64,\" + b64encode(mp4).decode()\n",
    "    display(\n",
    "        HTML(\n",
    "            \"\"\"\n",
    "          <video width=400 controls>\n",
    "                <source src=\"%s\" type=\"video/mp4\">\n",
    "          </video>\n",
    "      \"\"\"\n",
    "            % data_url\n",
    "        )\n",
    "    )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bddb2fce",
   "metadata": {},
   "outputs": [],
   "source": [
    "display_video(video_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd2aa926",
   "metadata": {},
   "outputs": [],
   "source": [
    "display_video(output_vid_trk_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2658e573",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "95a22eaad15357b8581a3e85dc7deb44b392b83f2a0c98c7572069d11874febe"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}