{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d7333de2-469c-4cf7-8dde-6b009204b3d4",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/anaconda3/envs/torch190-test/lib/python3.8/site-packages/tqdm/auto.py:22: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "import cv2\n",
    "import matplotlib.pyplot as plt\n",
    "from IPython.display import clear_output\n",
    "import os\n",
    "%matplotlib inline\n",
    "from PIL import Image\n",
    "import pandas as pd\n",
    "import time\n",
    "import natsort\n",
    "from tqdm import tqdm\n",
    "\n",
    "\n",
    "from fast_reid.fastreid.config import get_cfg\n",
    "from fast_reid.fastreid.modeling.meta_arch import build_model\n",
    "from fast_reid.fastreid.utils.checkpoint import Checkpointer\n",
    "from fast_reid.fastreid.engine import DefaultTrainer, default_argument_parser, default_setup, launch"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a879e3f8-4dc0-48d3-96f9-d50d30bd20c0",
   "metadata": {},
   "source": [
    "## Load obj in recent 3 frames"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "a91e710c-b6f1-4362-92a6-c42a2822ca5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# load obj_img and calculate average feat\n",
    "\n",
    "# average feat based matching\n",
    "\n",
    "# new id is in then get average\n",
    "\n",
    "tracking_result_path = '/data/cvprw/AIC23/tracking/segsort/run_demo/track_vis/botsort_val'\n",
    "\n",
    "channel_list = []\n",
    "temp = [file for file in os.listdir(tracking_result_path)]\n",
    "for channel in temp:\n",
    "    d = os.path.join(tracking_result_path, channel)\n",
    "    if os.path.isdir(d):\n",
    "#         print(channel)\n",
    "        channel_list.append(channel)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "3750d35c-d6e4-4f81-8842-932966fbf8b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "channel_list = natsort.natsorted(channel_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "a5890d94-9a92-4276-8555-da36e4f8095c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Open two video capture objects\n",
    "cap1 = cv2.VideoCapture(os.path.join(tracking_result_path, channel_list[0], 'video.mp4'))\n",
    "cap2 = cv2.VideoCapture(os.path.join(tracking_result_path, channel_list[1], 'video.mp4'))\n",
    "cap3 = cv2.VideoCapture(os.path.join(tracking_result_path, channel_list[2], 'video.mp4'))\n",
    "cap4 = cv2.VideoCapture(os.path.join(tracking_result_path, channel_list[3], 'video.mp4'))\n",
    "cap5 = cv2.VideoCapture(os.path.join(tracking_result_path, channel_list[4], 'video.mp4'))\n",
    "                        \n",
    "with open(os.path.join(tracking_result_path, channel_list[0], 'label.txt')) as label1:\n",
    "    gt1 = label1.readlines()\n",
    "with open(os.path.join(tracking_result_path, channel_list[1], 'label.txt')) as label2:\n",
    "    gt2 = label2.readlines()\n",
    "with open(os.path.join(tracking_result_path, channel_list[2], 'label.txt')) as label3:\n",
    "    gt3 = label3.readlines()\n",
    "with open(os.path.join(tracking_result_path, channel_list[3], 'label.txt')) as label4:\n",
    "    gt4 = label4.readlines()\n",
    "with open(os.path.join(tracking_result_path, channel_list[4], 'label.txt')) as label5:\n",
    "    gt5 = label5.readlines()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "4bc12dd6-8c53-409c-98ff-d54fec242ddf",
   "metadata": {},
   "outputs": [],
   "source": [
    "_, frame1 = cap1.read()\n",
    "_, frame2 = cap2.read()\n",
    "_, frame3 = cap3.read()\n",
    "_, frame4 = cap4.read()\n",
    "_, frame5 = cap5.read()\n",
    "\n",
    "width = cap1.get(cv2.CAP_PROP_FRAME_WIDTH)  # float\n",
    "height = cap1.get(cv2.CAP_PROP_FRAME_HEIGHT)  # float\n",
    "fps = cap1.get(cv2.CAP_PROP_FPS)\n",
    "vid_length = int(cap1.get(cv2.CAP_PROP_FRAME_COUNT))\n",
    "\n",
    "save_folder = '/data/cvprw/AIC23/dataset/GT_VID/val'\n",
    "\n",
    "vid_writer1 = cv2.VideoWriter(\n",
    "        os.path.join(save_folder, channel_list[0]+'.mp4'), cv2.VideoWriter_fourcc(*\"mp4v\"), fps, (int(width), int(height))\n",
    "    )\n",
    "\n",
    "vid_writer2 = cv2.VideoWriter(\n",
    "        os.path.join(save_folder, channel_list[1]+'.mp4'), cv2.VideoWriter_fourcc(*\"mp4v\"), fps, (int(width), int(height))\n",
    "    )\n",
    "\n",
    "vid_writer3 = cv2.VideoWriter(\n",
    "        os.path.join(save_folder, channel_list[2]+'.mp4'), cv2.VideoWriter_fourcc(*\"mp4v\"), fps, (int(width), int(height))\n",
    "    )\n",
    "\n",
    "vid_writer4 = cv2.VideoWriter(\n",
    "        os.path.join(save_folder, channel_list[3]+'.mp4'), cv2.VideoWriter_fourcc(*\"mp4v\"), fps, (int(width), int(height))\n",
    "    )\n",
    "\n",
    "vid_writer5 = cv2.VideoWriter(\n",
    "        os.path.join(save_folder, channel_list[4]+'.mp4'), cv2.VideoWriter_fourcc(*\"mp4v\"), fps, (int(width), int(height))\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "07147959-c712-46f8-9954-d89f09040007",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['2,33,690,294,36,131,1,-1,-1,-1\\n', '2,34,1734,484,80,224,1,-1,-1,-1\\n', '2,36,1258,282,26,109,1,-1,-1,-1\\n', '2,37,777,550,88,269,1,-1,-1,-1\\n', '2,38,506,465,83,233,1,-1,-1,-1\\n', '2,40,1523,317,57,109,1,-1,-1,-1\\n', '3,33,691,294,35,131,1,-1,-1,-1\\n', '3,34,1737,483,76,225,1,-1,-1,-1\\n', '3,37,778,550,88,270,1,-1,-1,-1\\n', '3,38,506,465,81,233,1,-1,-1,-1\\n']\n"
     ]
    }
   ],
   "source": [
    "print(gt1[:10])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "a5b38145-1752-4b3a-8f76-195295184fe0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def draw_bbox_tracking(img, results, thres):\n",
    "    for obj in results:\n",
    "        if obj[4]>thres:\n",
    "            x = int(obj[0])\n",
    "            y = int(obj[1])\n",
    "            x2 = int(obj[2])\n",
    "            y2 = int(obj[3])\n",
    "            \n",
    "            tracking_id = obj[5]\n",
    "            \n",
    "            color = (255,0,0)\n",
    "\n",
    "            tk = 1\n",
    "            # if obj['risk1'] == 'danger':\n",
    "            #     tk = 5\n",
    "            \n",
    "            content =  str(obj[5])+\"_\"+str(round(obj[4],2))\n",
    "            \n",
    "            font =  cv2.FONT_HERSHEY_PLAIN\n",
    "            img = cv2.rectangle(img, (x,y), (x2,y2), color, tk) # bbox\n",
    "            img = cv2.putText(img, content, (x, y-2), font, tk, color, tk, cv2.LINE_AA) # label\n",
    "    return img"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "eedcd5e4-48eb-47f6-bfd0-3b06de9d914c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract_tracking_result(gt, frame_count):\n",
    "    track_bbox=[]\n",
    "    for row in gt:\n",
    "        if int(row.split(',')[0]) == frame_count:\n",
    "            track_id=int(row.split(',')[1])\n",
    "            \n",
    "            x=int(row.split(',')[2])\n",
    "            y=int(row.split(',')[3])\n",
    "            w=int(row.split(',')[4])\n",
    "            h=int(row.split(',')[5])\n",
    "            \n",
    "            track_bbox.append([x,y,x+w,y+h,1,track_id])\n",
    "            \n",
    "        elif int(row.split(',')[0]) > frame_count:\n",
    "            break\n",
    "            \n",
    "    return track_bbox"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "d1c791f9-a783-45bb-a1af-e76d4fd77220",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 18010/18010 [1:01:03<00:00,  4.92it/s]\n"
     ]
    }
   ],
   "source": [
    "frame_count = 0\n",
    "\n",
    "t = tqdm(total=vid_length)\n",
    "\n",
    "while frame1 is not None:\n",
    "    track_bbox1 = extract_tracking_result(gt1, frame_count)\n",
    "    track_bbox2 = extract_tracking_result(gt2, frame_count)\n",
    "    track_bbox3 = extract_tracking_result(gt3, frame_count)\n",
    "    track_bbox4 = extract_tracking_result(gt4, frame_count)\n",
    "    track_bbox5 = extract_tracking_result(gt5, frame_count)\n",
    "    \n",
    "    result_img1 = draw_bbox_tracking(frame1, track_bbox1, 0.1)\n",
    "    result_img2 = draw_bbox_tracking(frame2, track_bbox1, 0.1)\n",
    "    result_img3 = draw_bbox_tracking(frame3, track_bbox1, 0.1)\n",
    "    result_img4 = draw_bbox_tracking(frame4, track_bbox1, 0.1)\n",
    "    result_img5 = draw_bbox_tracking(frame5, track_bbox1, 0.1)\n",
    "    \n",
    "    vid_writer1.write(result_img1)\n",
    "    vid_writer2.write(result_img2)\n",
    "    vid_writer3.write(result_img3)\n",
    "    vid_writer4.write(result_img4)\n",
    "    vid_writer5.write(result_img5)\n",
    "    \n",
    "    # plt.imshow(frame1)\n",
    "    # plt.show()\n",
    "    frame_count+=1\n",
    "    t.update(1)\n",
    "    \n",
    "    _, frame1 = cap1.read()\n",
    "    _, frame2 = cap2.read()\n",
    "    _, frame3 = cap3.read()\n",
    "    _, frame4 = cap4.read()\n",
    "    _, frame5 = cap5.read()\n",
    "    \n",
    "vid_writer1.release()\n",
    "vid_writer2.release()\n",
    "vid_writer3.release()\n",
    "vid_writer4.release()\n",
    "vid_writer5.release()\n",
    "\n",
    "t.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f3dbb67-d115-40d0-939d-7fc459ea97e7",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "torch190-test",
   "language": "python",
   "name": "torch190-test"
  },
  "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.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}