Towards-Realtime-MOT/live_track.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "d1489f9f-328c-4812-9cdb-0a2dee44ae88",
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "import argparse\n",
    "import logging\n",
    "import os\n",
    "from tqdm.auto import tqdm\n",
    "from pathlib import Path\n",
    "\n",
    "import numpy as np\n",
    "from tracker.multitracker import JDETracker\n",
    "from utils.timer import Timer\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from tracker.basetrack import TrackState\n",
    "import torch\n",
    "import pickle\n",
    "\n",
    "from utils import visualization as vis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "edec1b34-64ad-4610-856a-68d886a45142",
   "metadata": {},
   "outputs": [],
   "source": [
    "logger = logging.getLogger(__name__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "010bf567-8845-46d4-8500-883efce2d010",
   "metadata": {},
   "outputs": [],
   "source": [
    "# path of video to use. Replace with integer to use /dev/video device\n",
    "# that is, the webcam.\n",
    "video_path = Path('OUT/embedding_test/V003.seq')\n",
    "projector_pcl = Path('OUT/embedding_test/reducer_pca.pcl')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "8a413424-13c4-4bdc-825a-0aa6164e89e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils.parse_config import parse_model_cfg\n",
    "from utils.utils import mkdir_if_missing\n",
    "import utils.datasets as datasets\n",
    "from utils.log import logger as trmlog # we need to override this...\n",
    "\n",
    "\n",
    "trmlog.setLevel(logging.INFO)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "7b291b67-93ad-4b51-934a-dbaf095f7704",
   "metadata": {},
   "outputs": [],
   "source": [
    "# load options; quick'n'dirty copy from track.py (as the Namespace object is used in the multitracker)\n",
    "parser = argparse.ArgumentParser(prog='visualise_embeddings.py')\n",
    "parser.add_argument('--cfg', type=str, default='cfg/yolov3_1088x608.cfg', help='cfg file path')\n",
    "parser.add_argument('--weights', type=str, default='jde.1088x608.uncertainty.pt', help='path to weights file')\n",
    "parser.add_argument('--iou-thres', type=float, default=0.5, help='iou threshold required to qualify as detected')\n",
    "parser.add_argument('--conf-thres', type=float, default=0.5, help='object confidence threshold')\n",
    "parser.add_argument('--nms-thres', type=float, default=0.4, help='iou threshold for non-maximum suppression')\n",
    "parser.add_argument('--min-box-area', type=float, default=200, help='filter out tiny boxes')\n",
    "parser.add_argument('--track-buffer', type=int, default=30, help='tracking buffer')\n",
    "parser.add_argument('--dataset-dir', type=str, default=\"/datasets\", help='Path to directory with datasets')\n",
    "parser.add_argument('--experiment-name', type=str, default=\"embedding_test\", help=\"name to prefix output directory with\")\n",
    "parser.add_argument('--output-dir', type=str, default=\"./OUT\", help=\"directory for results\")\n",
    "    \n",
    "# we're running in notebook, so default to empty\n",
    "opt = parser.parse_args([])\n",
    "\n",
    "logger.setLevel(logging.INFO)\n",
    "result_path = os.path.join(opt.output_dir, opt.experiment_name)\n",
    "result_frame_path = os.path.join(result_path, 'track-test')\n",
    "mkdir_if_missing(result_path)\n",
    "mkdir_if_missing(result_frame_path)\n",
    "data_type = 'mot'\n",
    "\n",
    "# Read config\n",
    "cfg_dict = parse_model_cfg(opt.cfg)\n",
    "# set img_size in opt, so it is passed on to JDETracker\n",
    "opt.img_size = [int(cfg_dict[0]['width']), int(cfg_dict[0]['height'])]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "648faf4b-d692-473a-a99d-06b50a2e2261",
   "metadata": {},
   "outputs": [],
   "source": [
    "import queue, threading\n",
    "\n",
    "# bufferless VideoCapture (by Ulrich Stern: https://stackoverflow.com/a/54577746)\n",
    "# Usefull when using Webcam with higher fps than the script can process it.\n",
    "class BufferlessVideoCapture:\n",
    "\n",
    "  def __init__(self, name):\n",
    "    self.cap = cv2.VideoCapture(name)\n",
    "    self.q = queue.Queue()\n",
    "    t = threading.Thread(target=self._reader)\n",
    "    t.daemon = True\n",
    "    t.start()\n",
    "\n",
    "  # read frames as soon as they are available, keeping only most recent one\n",
    "  def _reader(self):\n",
    "    while True:\n",
    "      ret, frame = self.cap.read()\n",
    "      if not ret:\n",
    "        break\n",
    "      if not self.q.empty():\n",
    "        try:\n",
    "          self.q.get_nowait()   # discard previous (unprocessed) frame\n",
    "        except queue.Empty:\n",
    "          pass\n",
    "      self.q.put(frame)\n",
    "\n",
    "  def read(self):\n",
    "    return self.q.get()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "a28ef404-2031-43cf-aeb1-357aa1be0934",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "with projector_pcl.open('rb') as fp:\n",
    "   reducer = pickle.load(fp)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c433bdc3-e3bb-4d53-8e41-6e11791d8ac2",
   "metadata": {},
   "source": [
    "Load video file and get it's properties. Use that to calculate the dimension to fit the loaded model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "642c1953-6f33-4a22-a223-797755ece204",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "video 640 x 480\n",
      "Scale to 810 x 608\n"
     ]
    }
   ],
   "source": [
    "stream = cv2.VideoCapture(str(video_path))\n",
    "frame_rate = stream.get(cv2.CAP_PROP_FPS)\n",
    "vw = int(stream.get(cv2.CAP_PROP_FRAME_WIDTH))\n",
    "vh = int(stream.get(cv2.CAP_PROP_FRAME_HEIGHT))\n",
    "\n",
    "\n",
    "print(\"video\", vw, \"x\", vh)\n",
    "\n",
    "# get_size (from datasets.py)\n",
    "aspect_w, aspect_h = float(opt.img_size[0]) / vw, float(opt.img_size[1]) / vh\n",
    "aspect = min(aspect_w, aspect_h)\n",
    "w, h = int(vw *aspect), int(vh*aspect)\n",
    "print(\"Scale to\", w, \"x\", h)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "id": "be84553e-4309-46e5-9d9a-3075c20c0463",
   "metadata": {},
   "outputs": [],
   "source": [
    "umap_pcl = os.path.join(result_path, 'reducer_umap.pcl')\n",
    "pca_pcl = os.path.join(result_path, 'reducer_pca.pcl')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "id": "c623aa17-5ce2-4948-9adf-d4c9a6d1ccd2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(-1.0, 1.0)"
      ]
     },
     "execution_count": 61,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAABpQAAANkCAYAAAC588kTAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/bCgiHAAAACXBIWXMAAB7CAAAewgFu0HU+AACZQklEQVR4nOz9e5jXdZ0//t+R48DQ4AGLZtBNSA6eVygQE+lgfRr7eKqktfUQqO22hX2gNGvXpc9ubAu0UmrlIdRskw4eSCzNGnEhCbBL7ROHFdIV0AsDARkcjr5/f/Dj/RVhXswMMyJwu13X47qezOv5fj4fb50/3u/rPq/nq12SUgAAAAAAAKARh+zrBgAAAAAAAHhrEygBAAAAAABQSKAEAAAAAABAIYESAAAAAAAAhQRKAAAAAAAAFBIoAQAAAAAAUEigBAAAAAAAQCGBEgAAAAAAAIUESgAAAAAAABQSKAEAAAAAAFBIoAQAAAAAAEAhgRIAAAAAAACFBEoAAAAAAAAUEigBAAAAAABQSKAEAAAAAABAIYESAAAAAAAAhQRKAAAAAAAAFBIoAQAAAAAAUEigBAAAAAAAQCGBEgAAAAAAAIXaNFDq2bNnamtrM378+Dz44IP5y1/+klKplFKplKlTp7bJniNHjsxDDz2UF198MQ0NDXnuuefywx/+MEOGDGmT/QAAgIPT/vp9p6KiIl/60pcyd+7crF69OvX19Vm4cGEmTZqUo446qk36BgAADgyltqoiU6dObdW9unTpUnrggQca3W/r1q2lf/qnf2qz96qUUkoppZQ6uGp//L7Tp0+f0uLFixtdZ+3ataXa2tp9/t9WKaWUUkop9darN+3Iu//5n//JQw891Gbr/+AHP0htbW2S5Le//W3OOeecDB48OJ/5zGeyZMmStG/fPuPHj8/ll1/eZj0AAAAHp/3h+05lZWVmzJiRY489Nkly88035/3vf3+GDh2aa6+9NuvXr09VVVWmTZuWk046qc3eCwAAsP9qs7Tqn//5n0u1tbWlI488spSkdPTRR7fJX+yNGDGivO79999fOuSQQ3a6fvjhh5eee+65UqlUKr388sulHj167PMkTymllFJKKbV/1/72fWf8+PHldcaNG7fL9aFDh5Y2b95cKpVKpbq6un3+31cppZRSSin1lqs3b7O2+oI1Y8aMUqlUKm3evLlUXV292zkXXnhh4ZcnpZRSSimllNqbeit/3+nQoUNpzZo1pVKpVPrTn/5Uateu3W7X+e53v1teZ9CgQfv8v6lSSimllFLqrVNv2pF3baWysjIf+MAHkiSPPPJIVqxYsdt599xzT9atW5ckOe+88960/gAAAFqqtb7vjBgxIj169EiS3HHHHSmVSrtd5/bbby+PfW8CAABeb78PlAYPHpzOnTsnSWbOnNnovC1btmTOnDnl13To0OFN6Q8AAKClWuv7zumnn14eF60zf/78bNiwIUkybNiwFvcNAAAcePb7QGngwIHl8aJFiwrn7rjesWPHvPvd727TvgAAAPZWa33faeo627Zty5IlS5IkAwYMaHa/AADAgWu/v02npqamPF6+fHnh3GXLlpXHvXv3zsKFC5u8T3V1deH1Tp06ZcCAAXnppZfy0ksvZdu2bU1eGwAA9jft27fPkUcemSR5+umns3nz5n3c0YGptb7v7Finvr6+fDRe0TonnXRSjjzyyHTq1KnJ/299ZwIAgJ0daN+b9vtAqXv37uVxfX194dwdRzck288ib449fXkDAICD1eDBgzN//vx93cYBqbW+7+xYZ09r7G6dl19+uUm9+s4EAACNOxC+N+33R9516dKlPN5Turdp06byuKKios16AgAAaA2t9X1nxzpN+YtI35sAAIDd2e/vUNq4cWN53KlTp8K5Ox5mmyQNDQ3N2uf1R000dv31D8F98cUXm7U+AADsT3r16pV58+YlSV566aV93M2Bq7W+7+xYZ09r7GmdIr4zAQDAzg607037faC0fv368nhPx9h169atPG7KUQ+vt2LFiibPffHFF5s1HwAA9meehdN2Wuv7zo51mnL0d0u/N/nOBAAAjTsQvjft90fevf6c7j39RVzv3r3L49c/sBYAAOCtqLW+7+xYp7KyMlVVVU1a56WXXtrvHxoMAAC0nv0+UFqwYEF53L9//8K5O65v2bIlzzzzTJv2BQAAsLda6/tOU9dp3759+vTpkyRZuHBhs/sFAAAOXPt9oDRv3rzyQ2OHDx/e6LyOHTtmyJAh5dds3br1TekPAACgpVrr+86sWbPK46J1Bg0aVD4Wb/bs2S3uGwAAOPDs94FSfX19fvOb3yRJPvjBD6a6unq3884///zy0Q733nvvm9YfAABAS7XW951HH300a9euTZJccsklje536aWXlse+NwEAAK/3lg+ULrnkkpRKpZRKpVx33XW7nTNp0qQk2/8q78Ybb8whh+z8tg4//PB885vfTJKsWbMmt956a9s2DQAA0ARv1vedLVu25Nvf/naSZODAgRk3btwuc4YMGZJRo0Yl2R5AzZ8/v+VvDAAAOOB0aMvFhw0blr59+5b/fcQRR5THffv23eUv4+64444W7VNXV5cf//jH+dSnPpVzzjknv/71r3P99dfnhRdeyAknnJCvfvWrOfroo5MkV199dfkv8wAAAFpqf/u+M3HixFx44YXp169fJk6cmL59++buu+9OQ0NDRowYkWuvvTYdO3bMq6++mquuuqpFvQIAAAe2UlvV1KlTS82xuzUuueSS8vXrrruu0b26dOlSeuCBBxpde+vWrYWv39uqrq4u71VdXd1m+yillFJKKfVWKJ9/98/vO3369CktXry40XXWrl1bqq2t9TujlFJKKaVUK9SB9hn4LX/kXVNt3LgxZ599dv7mb/4mDz/8cFauXJlNmzbl+eefz49+9KOcfvrpGT9+/L5uEwAAoNla6/vO0qVLc8opp+TLX/5y5s2blzVr1mTDhg1ZtGhRvvWtb+XEE0/MjBkz3oR3BAAA7G/aZXuyxF6qrq7O8uXLkyQ1NTVZsWLFPu4IAADajs+/NJffGQAADjYH2mfgA+YOJQAAAAAAANqGQAkAAAAAAIBCAiUAAAAAAAAKCZQAAAAAAAAoJFACAAAAAACgkEAJAAAAAACAQgIlAAAAAAAACgmUAAAAAAAAKCRQAgAAAAAAoJBACQAAAAAAgEICJQAAAAAAAAoJlAAAAAAAACgkUAIAAAAAAKCQQAkAAAAAAIBCAiUAAAAAAAAKCZQAAAAAAAAoJFACAAAAAACgkEAJAAAAAACAQgIlAAAAAAAACgmUAAAAAAAAKCRQAgAAAAAAoJBACQAAAAAAgEICJQAAAAAAAAoJlAAAAAAAACgkUAIAAAAAAKCQQAkAAAAAAIBCAiUAAAAAAAAKCZQAAAAAAAAoJFACAAAAAACgkEAJAAAAAACAQgIlAAAAAAAACgmUAAAAAAAAKCRQAgAAAAAAoJBACQAAAAAAgEICJQAAAAAAAAoJlAAAAAAAACgkUAIAAAAAAKCQQAkAAAAAAIBCAiUAAAAAAAAKCZQAAAAAAAAoJFACAAAAAACgkEAJAAAAAACAQgIlAAAAAAAACgmUAAAAAAAAKCRQAgAAAAAAoJBACQAAAAAAgEICJQAAAAAAAAoJlAAAAAAAACgkUAIAAAAAAKCQQAkAAAAAAIBCAiUAAAAAAAAKCZQAAAAAAAAoJFACAAAAAACgkEAJAAAAAACAQgIlAAAAAAAACgmUAAAAAAAAKCRQAgAAAAAAoJBACQAAAAAAgEICJQAAAAAAAAoJlAAAAAAAACgkUAIAAAAAAKCQQAkAAAAAAIBCAiUAAAAAAAAKCZQAAAAAAAAoJFACAAAAAACgkEAJAAAAAACAQgIlAAAAAAAACgmUAAAAAAAAKCRQAgAAAAAAoJBACQAAAAAAgEICJQAAAAAAAAoJlAAAAAAAACgkUAIAAAAAAKCQQAkAAAAAAIBCAiUAAIC9dNRRR2XSpElZuHBh6uvrs3r16sydOzfjxo1LRUVFi9c9+uijUyqVmlXPPvvsbteqq6tr8hoAAAB
      "text/plain": [
       "<Figure size 2000x1000 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# matplotlib in interactive mode so the scatter can be updated.\n",
    "\n",
    "plt.style.use('dark_background')\n",
    "\n",
    "fig, ax = plt.subplots(1, 2, figsize=(10, 5), dpi=200)\n",
    "x, y = [],[]\n",
    "scatter = ax[1].scatter(x,y)\n",
    "# scatter = ax.plot([],[],label='toto',ms=10,color='k',marker='o',ls='')\n",
    "plt.xlim(-1,1)\n",
    "plt.ylim(-1,1)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "id": "e8952235-7e56-4606-858a-a9165b967726",
   "metadata": {},
   "outputs": [],
   "source": [
    "# init tracker\n",
    "tracker = JDETracker(opt, frame_rate=frame_rate)\n",
    "timer = Timer()\n",
    "results = []\n",
    "frame_id = -1\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9cc6fd1-b9c2-4303-a21c-a193c6045526",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3ec1f9b647cf40edbf662db14ac36bcb",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "with tqdm() as pbar:\n",
    "    while True:\n",
    "        timer.tic()\n",
    "        frame_id += 1\n",
    "\n",
    "        pbar.update(1)\n",
    "\n",
    "        ret, frame = stream.read()\n",
    "        # scale down/up frame to fit tracker\n",
    "        frame = cv2.resize(frame, (w, h))\n",
    "\n",
    "\n",
    "        # letterbox as to have a constant size for the model\n",
    "        img, _, _, _ = datasets.letterbox(frame, height=opt.img_size[1], width=opt.img_size[0])\n",
    "\n",
    "        # Normalize RGB\n",
    "        img = img[:, :, ::-1].transpose(2, 0, 1)\n",
    "        img = np.ascontiguousarray(img, dtype=np.float32)\n",
    "        img /= 255.0\n",
    "\n",
    "        # send frame to GPU\n",
    "        blob = torch.from_numpy(img).cuda().unsqueeze(0)\n",
    "\n",
    "        # online targets: all targets that are not timed out\n",
    "        # frame_embeddings: the embeddings of objects visible only in the current frame\n",
    "        online_targets, frame_embeddings = tracker.update(blob, frame)\n",
    "\n",
    "        # all relevant tracks from the tracker directly (excludes the 'removed' ones)\n",
    "        # thus also accessing its those that it has stored from previous frames\n",
    "        all_stracks = tracker.tracked_stracks + tracker.lost_stracks\n",
    "\n",
    "        # project them to 2D using PCA (see visualise_embeddings.ipynb)\n",
    "        if len(all_stracks):\n",
    "            projections = reducer.transform([strack.smooth_feat for strack in all_stracks])\n",
    "        else:\n",
    "            projections = []\n",
    "\n",
    "        timer.toc()\n",
    "\n",
    "\n",
    "        online_tlwhs = []\n",
    "        online_ids = []\n",
    "        for t in online_targets:\n",
    "            tlwh = t.tlwh\n",
    "            tid = t.track_id\n",
    "            vertical = tlwh[2] / tlwh[3] > 1.6\n",
    "            if tlwh[2] * tlwh[3] > opt.min_box_area and not vertical:\n",
    "                online_tlwhs.append(tlwh)\n",
    "                online_ids.append(tid)\n",
    "        # visualise results for single frame\n",
    "        online_im = vis.plot_tracking(frame, online_tlwhs, online_ids, frame_id=frame_id,\n",
    "                                            fps=1. / timer.average_time)\n",
    "        \n",
    "        # redraw the canvas\n",
    "        # fig.canvas.draw()\n",
    "        # ax[0].clear()\n",
    "        # ax[1].clear()\n",
    "        \n",
    "        ax[0].imshow(online_im)\n",
    "        # cv2.imwrite(result_frame_path + f\"/{frame_id:04d}-live.jpg\", online_im)\n",
    "\n",
    "\n",
    "        # scatter.set_data(projections[:,0],projections[:,1])\n",
    "        if len(projections):\n",
    "            scatter.set_offsets(projections)\n",
    "            scatter.set_color([(0,0,1,0.5) if strack.state == TrackState.Lost else (0,0,1,1) for strack in all_stracks])\n",
    "        else:\n",
    "            scatter = ax[1].scatter([],[])\n",
    "            pbar.set_description('todo: empty axes')\n",
    "            # print()\n",
    "\n",
    "    #     plot = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep='')\n",
    "    #     plot  = plot.reshape(fig.canvas.get_width_height()[::-1] + (3,))\n",
    "\n",
    "    #     # img is rgb, convert to opencv's default bgr\n",
    "    #     plot = cv2.cvtColor(plot,cv2.COLOR_RGB2BGR)\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "        fig.savefig(result_frame_path + f\"/{frame_id:04d}.png\")\n",
    "\n",
    "        # fig.canvas.draw_idle()\n",
    "        # plt.pause(0.1)\n",
    "    #     if cv2.waitKey(1) & 0xFF == ord('q'): # wait for 1 millisecond\n",
    "    #         break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f21e642-2320-4a7e-a0ea-af32e8d1a182",
   "metadata": {},
   "outputs": [],
   "source": [
    "! ffmpeg -i OUT/embedding_test/track-test/%04d.png -c:v libx264 -crf 10  OUT/embedding_test/track-test.mp4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94d6e0cb-4d32-4d18-8b18-4ecfcc67a634",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30953716-c267-401f-ad30-51b2a139bfc2",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "trm",
   "language": "python",
   "name": "trm"
  },
  "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.10.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}