{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d1fa6b91",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/anaconda3/envs/torch190-test/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\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": "674e8d49",
   "metadata": {},
   "source": [
    "## Load obj in recent 3 frames"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "1945922c",
   "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/SUBMISSION/S018'\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": "8baa6c4a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['.ipynb_checkpoints', 'c100', 'c101', 'c102', 'c103', 'c104', 'c105']"
      ]
     },
     "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": "50a0d9d6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['c100', 'c101', 'c102', 'c103', 'c104', 'c105']"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "del channel_list[0]\n",
    "channel_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "c7c561fb",
   "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('/data/cvprw/AIC23/tracking/paper_vis/Trim.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], 'bot_sort.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('/data/cvprw/AIC23/tracking/paper_vis/Byte.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": 11,
   "id": "253b5843",
   "metadata": {},
   "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 = cap5.get(cv2.CAP_PROP_FRAME_WIDTH)  # float\n",
    "height = cap5.get(cv2.CAP_PROP_FRAME_HEIGHT)  # float\n",
    "fps = cap5.get(cv2.CAP_PROP_FPS)\n",
    "vid_length = int(cap5.get(cv2.CAP_PROP_FRAME_COUNT))\n",
    "\n",
    "save_folder = '/data/cvprw/AIC23/tracking/paper_vis/result'\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, 'Byte.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": 12,
   "id": "f079d460",
   "metadata": {},
   "outputs": [],
   "source": [
    "def draw_bbox_tracking(img, results, thres,frame_count,tracklet_history):\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",
    "            center_x = int((x+x2)/2)\n",
    "            # Update history\n",
    "            if tracking_id not in tracklet_history:\n",
    "                tracklet_history[tracking_id] = []\n",
    "            tracklet_history[tracking_id].append((center_x, y2))\n",
    "            \n",
    "            color = give_color(tracking_id%100)\n",
    "            \n",
    "            # Draw historical path\n",
    "#             if len(tracklet_history[tracking_id]) > 1:\n",
    "#                 for i in range(1, len(tracklet_history[tracking_id])):\n",
    "#                     cv2.line(img, tracklet_history[tracking_id][i - 1], tracklet_history[tracking_id][i], color, 3)\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, tracklet_history"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "506f1856",
   "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(float(row.split(',')[2]))\n",
    "            y=int(float(row.split(',')[3]))\n",
    "            w=int(float(row.split(',')[4]))\n",
    "            h=int(float(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": 14,
   "id": "e020c4ed",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████| 212/212 [00:02<00:00, 78.14it/s]\n"
     ]
    }
   ],
   "source": [
    "frame_count = 0\n",
    "tracklet_history = {}\n",
    "\n",
    "t = tqdm(total=vid_length)\n",
    "\n",
    "while frame5 is not None and frame_count<300:\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,tracklet_history  = draw_bbox_tracking(frame1, track_bbox1, 0.1,frame_count,tracklet_history)\n",
    "    result_img5,tracklet_history  = draw_bbox_tracking(frame5, track_bbox5, 0.1,frame_count,tracklet_history)\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": "3864acbe",
   "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
}