{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "9d6ec69e",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/anaconda3/envs/scit/lib/python3.8/site-packages/tqdm/auto.py:21: 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",
    "from multicam_utils import give_color"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8727303b",
   "metadata": {},
   "source": [
    "## Load obj in recent 3 frames"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "4ef63fd1",
   "metadata": {
    "tags": []
   },
   "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/SUBMISSION/S022'\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": "b2e2e3bd",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['.ipynb_checkpoints', 'c124', 'c125', 'c126', 'c127', 'c128', 'c129']"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "channel_list = natsort.natsorted(channel_list)\n",
    "channel_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "bb916615",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['c124', 'c125', 'c126', 'c127', 'c128', 'c129']"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "del channel_list[0]\n",
    "channel_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "2b4af650",
   "metadata": {
    "tags": []
   },
   "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",
    "cap6 = cv2.VideoCapture(os.path.join(tracking_result_path, channel_list[5], 'video.mp4'))\n",
    "                        \n",
    "with open(os.path.join(tracking_result_path, channel_list[0], 'sorted_associated.txt')) as label1:\n",
    "    gt1 = label1.readlines()\n",
    "with open(os.path.join(tracking_result_path, channel_list[1], 'sorted_associated.txt')) as label2:\n",
    "    gt2 = label2.readlines()\n",
    "with open(os.path.join(tracking_result_path, channel_list[2], 'sorted_associated.txt')) as label3:\n",
    "    gt3 = label3.readlines()\n",
    "with open(os.path.join(tracking_result_path, channel_list[3], 'sorted_associated.txt')) as label4:\n",
    "    gt4 = label4.readlines()\n",
    "with open(os.path.join(tracking_result_path, channel_list[4], 'sorted_associated.txt')) as label5:\n",
    "    gt5 = label5.readlines()\n",
    "with open(os.path.join(tracking_result_path, channel_list[5], 'sorted_associated.txt')) as label6:\n",
    "    gt6 = label6.readlines()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "a1bf7fa8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "_, frame1 = cap1.read()\n",
    "_, frame2 = cap2.read()\n",
    "_, frame3 = cap3.read()\n",
    "_, frame4 = cap4.read()\n",
    "_, frame5 = cap5.read()\n",
    "_, frame6 = cap6.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/tracking/SUBMISSION/S022'\n",
    "\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",
    "    )\n",
    "\n",
    "vid_writer6 = cv2.VideoWriter(\n",
    "        os.path.join(save_folder, channel_list[5]+'.mp4'), cv2.VideoWriter_fourcc(*\"mp4v\"), fps, (int(width), int(height))\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "889fa98e",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['0,103,0,1327,0,79,148,-1,-1\\n', '0,100,0,778,212,123,401,1258,436\\n', '0,104,0,555,869,210,210,-1,-1\\n', '0,100,1,778,213,122,399,1258,436\\n', '0,104,1,554,869,211,210,-1,-1\\n', '0,103,1,1328,0,78,148,-1,-1\\n', '0,103,2,1330,0,77,147,-1,-1\\n', '0,104,2,551,871,214,208,-1,-1\\n', '0,100,2,778,213,122,398,1257,436\\n', '0,100,3,778,213,118,398,1257,436\\n']\n"
     ]
    }
   ],
   "source": [
    "print(gt1[:10])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "d53c4c58",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def draw_bbox_tracking(img, results, thres,frame_count):\n",
    "    font =  cv2.FONT_HERSHEY_PLAIN\n",
    "    \n",
    "    img = cv2.putText(img, \"frame count : \"+str(frame_count), (30, 30), font, 3, (255,0,0), 3, cv2.LINE_AA) # label\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 = give_color(tracking_id%100)\n",
    "\n",
    "            tk = 3\n",
    "            # if obj['risk1'] == 'danger':\n",
    "            #     tk = 5\n",
    "            \n",
    "            content = str(tracking_id)\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": 9,
   "id": "4e86f3c2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def extract_tracking_result(gt, frame_count):\n",
    "    track_bbox=[]\n",
    "    for row in gt:\n",
    "        if int(row.split(',')[2]) == frame_count:\n",
    "            track_id=int(row.split(',')[1])\n",
    "            \n",
    "            x=int(row.split(',')[3])\n",
    "            y=int(row.split(',')[4])\n",
    "            w=int(row.split(',')[5])\n",
    "            h=int(row.split(',')[6])\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": 10,
   "id": "f1dac293",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  3%|███▋                                                                                                                                  | 500/18010 [00:41<24:28, 11.93it/s]\n"
     ]
    }
   ],
   "source": [
    "frame_count = 0\n",
    "\n",
    "t = tqdm(total=vid_length)\n",
    "\n",
    "while frame1 is not None and frame_count<500:\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",
    "    track_bbox6 = extract_tracking_result(gt6, frame_count)\n",
    "    \n",
    "    result_img1 = draw_bbox_tracking(frame1, track_bbox1, 0.1,frame_count)\n",
    "    result_img2 = draw_bbox_tracking(frame2, track_bbox2, 0.1,frame_count)\n",
    "    result_img3 = draw_bbox_tracking(frame3, track_bbox3, 0.1,frame_count)\n",
    "    result_img4 = draw_bbox_tracking(frame4, track_bbox4, 0.1,frame_count)\n",
    "    result_img5 = draw_bbox_tracking(frame5, track_bbox5, 0.1,frame_count)\n",
    "    result_img6 = draw_bbox_tracking(frame6, track_bbox6, 0.1,frame_count)\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",
    "    vid_writer6.write(result_img6)\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",
    "    _, frame6 = cap6.read()\n",
    "    \n",
    "vid_writer1.release()\n",
    "vid_writer2.release()\n",
    "vid_writer3.release()\n",
    "vid_writer4.release()\n",
    "vid_writer5.release()\n",
    "vid_writer6.release()\n",
    "\n",
    "t.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fce4f858",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "scit",
   "language": "python",
   "name": "scit"
  },
  "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
}