{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "5e46d2e1-d240-4e0f-bf12-d5dd67cd0c56",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/anaconda3/envs/mmyolo/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 argparse\n",
    "import os\n",
    "import os.path as osp\n",
    "import time\n",
    "import cv2\n",
    "import torch\n",
    "import numpy as np\n",
    "from loguru import logger\n",
    "from collections import defaultdict, Counter\n",
    "\n",
    "from mmdet.apis import async_inference_detector, inference_detector\n",
    "from mmdet.apis.inference import init_detector\n",
    "\n",
    "\n",
    "from yolox.utils.visualize import plot_tracking, _COLORS\n",
    "from yolox.tracker.tp_tracker_interpolation import TPTracker\n",
    "from yolox.tracking_utils.timer import Timer\n",
    "\n",
    "from model import SocialImplicit\n",
    "from CFG import CFG\n",
    "import numpy as np\n",
    "from scipy.fftpack import fft, ifft\n",
    "from metrics import *\n",
    "\n",
    "IMAGE_EXT = [\".jpg\", \".jpeg\", \".webp\", \".bmp\", \".png\"]\n",
    "\n",
    "from IPython.display import clear_output\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "# clear_output(wait=False)\n",
    "\n",
    "from IPython.display import display, Javascript\n",
    "\n",
    "def clear_all_output():\n",
    "    display(Javascript('IPython.notebook.clear_all_output()'))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49a53d22-2616-4d9b-b67c-e044eb09c8a4",
   "metadata": {},
   "source": [
    "# 1. Single-Cam Trajectory Prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "07a05393-5d44-4d46-af3c-f4865c01eced",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def make_parser():\n",
    "    parser = argparse.ArgumentParser(\"Trajectory-Prediction Track Demo!\")\n",
    "    parser.add_argument(\n",
    "        \"demo\", default=\"video\", help=\"demo type, eg. image, video and webcam\"\n",
    "    )\n",
    "    \n",
    "    ########### mmdet bbox detector ##############\n",
    "    parser.add_argument('video', help='Video file path')\n",
    "    # parser.add_argument('config', help='Config file')\n",
    "    # parser.add_argument('checkpoint', help='Checkpoint file')\n",
    "    parser.add_argument(\n",
    "        '--device', default='cuda:0', help='Device used for inference')\n",
    "    ##############################################\n",
    "    \n",
    "    ########### trajectory predictor ##############\n",
    "    parser.add_argument('--tp_weight', type=str, help='trajectory prediction model weight')\n",
    "    ##############################################\n",
    "\n",
    "    parser.add_argument('--save_result', type=str, help='Output video file')\n",
    "    parser.add_argument('--save_vid', default=False, help='Output video file')\n",
    "    \n",
    "    # tracking args\n",
    "    parser.add_argument(\"--track_thresh\", type=float, default=0.8, help=\"tracking confidence threshold\")\n",
    "    parser.add_argument(\"--track_buffer\", type=int, default=150, help=\"the frames for keep lost tracks\")\n",
    "    parser.add_argument(\"--match_thresh\", type=float, default=0.8, help=\"matching threshold for tracking\")\n",
    "    # parser.add_argument(\n",
    "    #     \"--aspect_ratio_thresh\", type=float, default=1.6,\n",
    "    #     help=\"threshold for filtering out boxes of which aspect ratio are above the given value.\"\n",
    "    # )\n",
    "    # parser.add_argument('--min_box_area', type=float, default=10, help='filter out tiny boxes')\n",
    "    # parser.add_argument(\"--mot20\", dest=\"mot20\", default=False, action=\"store_true\", help=\"test mot20.\")\n",
    "    return parser\n",
    "\n",
    "args_list = [\n",
    "    'video',\n",
    "    '/data/autocon/hd_dataset/site1_093448',  # replace with actual path\n",
    "    # './weights/yolox_x.py',  # replace with actual path\n",
    "    # './weights/epoch_250.pth',  # replace with actual path\n",
    "    '--device', 'cuda:0',\n",
    "    '--tp_weight', './weights/tp_best.pth',  # replace with actual value\n",
    "    '--save_result', './site1_1-result',  # replace with actual path\n",
    "    '--save_vid', 'False',  # or 'False'\n",
    "    '--track_thresh', '0.8',\n",
    "    '--track_buffer', '30',\n",
    "    '--match_thresh', '0.8',\n",
    "    # '--min_box_area', '10',\n",
    "    # Add any other argument values you want to set\n",
    "]\n",
    "\n",
    "args = make_parser().parse_args(args_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "fd5700b3-41ee-476b-aef5-81efc8bc7d2b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def mmdet3x_convert(results, threshold):\n",
    "    det_result_list = []\n",
    "   \n",
    "    confidence_score = results.pred_instances.scores.tolist()\n",
    "    for i, conf in enumerate(confidence_score):\n",
    "        if conf >= threshold:\n",
    "            # print(conf)\n",
    "            \n",
    "            extracted_box = results.pred_instances.bboxes[i].cpu().tolist()\n",
    "            extracted_label = results.pred_instances.labels[i].cpu().tolist()\n",
    "            \n",
    "            class_name = 0\n",
    "            \n",
    "            temp = [float(extracted_box[0]),\n",
    "                     float(extracted_box[1]),\n",
    "                     float(extracted_box[2]),\n",
    "                     float(extracted_box[3]), float(0.9), class_name]\n",
    "            \n",
    "            det_result_list.append(temp)\n",
    "            \n",
    "    return np.array(det_result_list)\n",
    "\n",
    "def mmdet3x_convert_pretrain(results, threshold):\n",
    "    det_result_list = []\n",
    "\n",
    "    permit_class_list = [0, 7]\n",
    "    class_labels = results.pred_instances.labels.tolist()\n",
    "    \n",
    "    confidence_score = results.pred_instances.scores.tolist()\n",
    "    for i, conf in enumerate(confidence_score):\n",
    "        if conf >= threshold and class_labels[i] in permit_class_list:\n",
    "            # print(conf)\n",
    "            \n",
    "            extracted_box = results.pred_instances.bboxes[i].cpu().tolist()\n",
    "            extracted_label = results.pred_instances.labels[i].cpu().tolist()\n",
    "            \n",
    "            class_name = 0\n",
    "            \n",
    "            temp = [float(extracted_box[0]),\n",
    "                     float(extracted_box[1]),\n",
    "                     float(extracted_box[2]),\n",
    "                     float(extracted_box[3]), float(0.9), class_name]\n",
    "            \n",
    "            det_result_list.append(temp)\n",
    "            \n",
    "    return np.array(det_result_list)\n",
    "\n",
    "def compute_iou(boxA, boxB):\n",
    "    # Compute the coordinates of the intersection rectangle\n",
    "    xA = max(boxA[0], boxB[0])\n",
    "    yA = max(boxA[1], boxB[1])\n",
    "    xB = min(boxA[2], boxB[2])\n",
    "    yB = min(boxA[3], boxB[3])\n",
    "    \n",
    "    # Compute the area of intersection\n",
    "    interArea = max(0, xB - xA) * max(0, yB - yA)\n",
    "    \n",
    "    # Compute the area of both boxes\n",
    "    boxAArea = (boxA[2] - boxA[0]) * (boxA[3] - boxA[1])\n",
    "    boxBArea = (boxB[2] - boxB[0]) * (boxB[3] - boxB[1])\n",
    "    \n",
    "    # Compute IoU\n",
    "    iou = interArea / float(boxAArea + boxBArea - interArea)\n",
    "    \n",
    "    return iou\n",
    "\n",
    "def nms(boxes, overlap_threshold=0.5):\n",
    "    if len(boxes) == 0:\n",
    "        return []\n",
    "    \n",
    "    pick = []\n",
    "    \n",
    "    # Extract coordinates and scores from the boxes\n",
    "    x1 = boxes[:,0]\n",
    "    y1 = boxes[:,1]\n",
    "    x2 = boxes[:,2]\n",
    "    y2 = boxes[:,3]\n",
    "    scores = boxes[:,4]\n",
    "    \n",
    "    # Compute the area of each box and sort the scores in descending order\n",
    "    area = (x2 - x1) * (y2 - y1)\n",
    "    indices = np.argsort(scores)[::-1]\n",
    "    \n",
    "    while len(indices) > 0:\n",
    "        # Take the index of the box with the highest score\n",
    "        current = indices[0]\n",
    "        pick.append(current)\n",
    "        \n",
    "        # Compute IoU between the box with the highest score and the other boxes\n",
    "        ious = np.array([compute_iou(boxes[current], boxes[i]) for i in indices[1:]])\n",
    "        \n",
    "        # Get the indices of boxes that are overlapping less than the threshold with the current box\n",
    "        remove_indices = np.where(ious > overlap_threshold)[0] + 1\n",
    "        \n",
    "        # Remove the current box and the overlapping boxes\n",
    "        indices = np.delete(indices, remove_indices)\n",
    "        indices = np.delete(indices, 0)\n",
    "        \n",
    "    return boxes[pick]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "b0d3ed36-7b75-4495-a26e-9dccfeb1d2bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "frame_img_list = []\n",
    "traj_list = []\n",
    "\n",
    "def imageflow_demo(det_model, vis_folder, current_time, args, tp_model):\n",
    "    \n",
    "    tracker_list = []\n",
    "    \n",
    "    temp = []\n",
    "    channel_list = [file for file in os.listdir(args.video)]\n",
    "    for channel in channel_list:\n",
    "        d = os.path.join(args.video, channel)\n",
    "        if os.path.isdir(d) and not channel.startswith('.'):\n",
    "            temp.append(channel)\n",
    "    \n",
    "    for channel in temp:\n",
    "        video_path = os.path.join(os.path.join(args.video, channel, 'video.mp4'))\n",
    "    \n",
    "        cap = cv2.VideoCapture(video_path)\n",
    "        width = cap.get(cv2.CAP_PROP_FRAME_WIDTH)  # float\n",
    "        height = cap.get(cv2.CAP_PROP_FRAME_HEIGHT)  # float\n",
    "        fps = cap.get(cv2.CAP_PROP_FPS)\n",
    "        timestamp = time.strftime(\"%Y_%m_%d_%H_%M_%S\", current_time)\n",
    "        \n",
    "        if args.save_vid:\n",
    "            save_folder = osp.join(vis_folder, channel)\n",
    "            os.makedirs(save_folder, exist_ok=True)\n",
    "            save_path = osp.join(save_folder,\"inference_vid.mp4\")\n",
    "\n",
    "            logger.info(f\"video save_path is {save_path}\")\n",
    "            vid_writer = cv2.VideoWriter(\n",
    "                save_path, cv2.VideoWriter_fourcc(*\"mp4v\"), fps, (int(width), int(height))\n",
    "            )\n",
    "        \n",
    "        tracker = TPTracker(tp_model, args, frame_rate=30)\n",
    "        timer = Timer()\n",
    "        frame_id = 0\n",
    "        results = []\n",
    "        \n",
    "        while True:\n",
    "            if frame_id % 30 == 0:\n",
    "                logger.info('Processing frame {} ({:.2f} fps)'.format(frame_id, 1. / max(1e-5, timer.average_time)))\n",
    "            ret_val, frame = cap.read()\n",
    "            if ret_val:\n",
    "                \n",
    "                timer.tic()\n",
    "                worker_det_results = inference_detector(det_model[0], frame)\n",
    "                vehicle_det_results = inference_detector(det_model[1], frame)\n",
    "                pretrained_det_results = inference_detector(det_model[2], frame)\n",
    "                \n",
    "                threshold = 0.15\n",
    "                worker_class_dic = {0: 8, 1: 9}\n",
    "                vehicle_class_dic = {0 : 0, 1 : 1, 2 : 2, 3 : 3, 4 : 4, 5 : 5, 6 : 6, 7 : 7}\n",
    "                \n",
    "                worker_results = mmdet3x_convert(worker_det_results, threshold)\n",
    "                vehicle_results = mmdet3x_convert(vehicle_det_results, threshold)\n",
    "                pretrained_results = mmdet3x_convert_pretrain(pretrained_det_results, threshold)\n",
    "                \n",
    "                # *** dets format [x1,y1,x2,y2,score,class_id]\n",
    "                all_results = [result for result in [vehicle_results, worker_results, pretrained_results] if result.size > 0]\n",
    "\n",
    "                if len(all_results) == 0:\n",
    "                    final_boxes = np.array([], dtype=float64)\n",
    "                else:\n",
    "                    combined_boxes = np.concatenate(all_results, axis=0)\n",
    "                    final_boxes = nms(combined_boxes)\n",
    "\n",
    "                dets_arr = np.array(final_boxes)\n",
    "\n",
    "                # ##############\n",
    "                # print(\"frame # : \", frame_id)\n",
    "                # print(dets_arr)\n",
    "                # ##############\n",
    "              \n",
    "                if dets_arr is not None:\n",
    "                    # Run tracker\n",
    "                    online_targets = tracker.update(dets_arr, frame)\n",
    "                    \n",
    "                    online_tlwhs = []\n",
    "                    online_ids = []\n",
    "                    online_scores = []\n",
    "                    online_traj = []\n",
    "                    for t in online_targets:\n",
    "                        tlwh = t.tlwh\n",
    "                        tid = t.track_id\n",
    "                        # vertical = tlwh[2] / tlwh[3] > args.aspect_ratio_thresh\n",
    "                        online_tlwhs.append(tlwh)\n",
    "                        online_ids.append(tid)\n",
    "                        online_scores.append(t.score)\n",
    "                        online_traj.append(t.tlwh_queue.to_array())\n",
    "                        results.append(\n",
    "                            f\"{frame_id},{tid},{tlwh[0]:.2f},{tlwh[1]:.2f},{tlwh[2]:.2f},{tlwh[3]:.2f},{t.score:.2f},-1,-1,-1\\n\"\n",
    "                        )\n",
    "                    timer.toc()\n",
    "                    \n",
    "                    online_im = plot_tracking(\n",
    "                        frame, online_tlwhs, online_ids, online_traj, frame_id=frame_id + 1, fps=1. / timer.average_time\n",
    "                    )\n",
    "                else:\n",
    "                    timer.toc()\n",
    "                    online_im = frame\n",
    "                if args.save_vid:\n",
    "                    vid_writer.write(online_im)\n",
    "                ch = cv2.waitKey(1)\n",
    "                if ch == 27 or ch == ord(\"q\") or ch == ord(\"Q\"):\n",
    "                    break\n",
    "            else:\n",
    "                break\n",
    "                \n",
    "            # if frame_id == 51:\n",
    "            #     frame_img_list.append(online_im) \n",
    "            #     traj_list.append(online_traj)\n",
    "            #     break    \n",
    "            \n",
    "            frame_id += 1\n",
    "\n",
    "        if args.save_result:\n",
    "            res_file = osp.join(save_folder, f\"{timestamp}.txt\")\n",
    "            with open(res_file, 'w') as f:\n",
    "                f.writelines(results)\n",
    "            logger.info(f\"save results to {res_file}\")\n",
    "            \n",
    "        tracker_list.append(tracker)\n",
    "        \n",
    "    return tracker_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "bb5199b5-c0b0-4b5b-92e5-d421c160a149",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[32m2023-12-26 21:02:00.232\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36m<module>\u001b[0m:\u001b[36m23\u001b[0m - \u001b[1mArgs: Namespace(demo='video', device='cuda:0', match_thresh=0.8, save_result='./site1_1-result', save_vid='False', tp_weight='./weights/tp_best.pth', track_buffer=30, track_thresh=0.8, video='/data/autocon/hd_dataset/site1_093448')\u001b[0m\n",
      "\u001b[32m2023-12-26 21:02:00.318\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mimageflow_demo\u001b[0m:\u001b[36m29\u001b[0m - \u001b[1mvideo save_path is ./site1_1-result/track_vis/c5/inference_vid.mp4\u001b[0m\n",
      "\u001b[32m2023-12-26 21:02:00.322\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mimageflow_demo\u001b[0m:\u001b[36m41\u001b[0m - \u001b[1mProcessing frame 0 (100000.00 fps)\u001b[0m\n",
      "/opt/anaconda3/envs/mmyolo/lib/python3.8/site-packages/torch/functional.py:445: UserWarning: torch.meshgrid: in an upcoming release, it will be required to pass the indexing argument. (Triggered internally at  /opt/conda/conda-bld/pytorch_1639180588308/work/aten/src/ATen/native/TensorShape.cpp:2157.)\n",
      "  return _VF.meshgrid(tensors, **kwargs)  # type: ignore[attr-defined]\n",
      "\u001b[32m2023-12-26 21:02:04.265\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mimageflow_demo\u001b[0m:\u001b[36m41\u001b[0m - \u001b[1mProcessing frame 30 (9.12 fps)\u001b[0m\n",
      "/data/autocon/SCIT-MCMT-Tracking/yolox/tracker/tp_tracker_interpolation.py:223: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
      "  obs_traj = torch.tensor(obs_traj, device='cuda:0', dtype=torch.float64)\n",
      "/data/autocon/SCIT-MCMT-Tracking/yolox/tracker/tp_tracker_interpolation.py:224: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
      "  V_obs = torch.tensor(V_obs, device='cuda:0', dtype=torch.float64)\n",
      "\u001b[32m2023-12-26 21:02:07.387\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mimageflow_demo\u001b[0m:\u001b[36m41\u001b[0m - \u001b[1mProcessing frame 60 (10.19 fps)\u001b[0m\n",
      "\u001b[32m2023-12-26 21:02:10.816\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mimageflow_demo\u001b[0m:\u001b[36m41\u001b[0m - \u001b[1mProcessing frame 90 (10.24 fps)\u001b[0m\n",
      "\u001b[32m2023-12-26 21:02:13.979\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mimageflow_demo\u001b[0m:\u001b[36m41\u001b[0m - \u001b[1mProcessing frame 120 (10.49 fps)\u001b[0m\n",
      "\u001b[32m2023-12-26 21:02:14.099\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mimageflow_demo\u001b[0m:\u001b[36m120\u001b[0m - \u001b[1msave results to ./site1_1-result/track_vis/c5/2023_12_26_21_02_00.txt\u001b[0m\n",
      "\u001b[32m2023-12-26 21:02:14.164\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mimageflow_demo\u001b[0m:\u001b[36m29\u001b[0m - \u001b[1mvideo save_path is ./site1_1-result/track_vis/c3/inference_vid.mp4\u001b[0m\n",
      "\u001b[32m2023-12-26 21:02:14.275\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mimageflow_demo\u001b[0m:\u001b[36m41\u001b[0m - \u001b[1mProcessing frame 0 (100000.00 fps)\u001b[0m\n",
      "\u001b[32m2023-12-26 21:02:17.398\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mimageflow_demo\u001b[0m:\u001b[36m41\u001b[0m - \u001b[1mProcessing frame 30 (12.03 fps)\u001b[0m\n",
      "\u001b[32m2023-12-26 21:02:20.393\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mimageflow_demo\u001b[0m:\u001b[36m41\u001b[0m - \u001b[1mProcessing frame 60 (12.01 fps)\u001b[0m\n",
      "\u001b[32m2023-12-26 21:02:23.459\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mimageflow_demo\u001b[0m:\u001b[36m41\u001b[0m - \u001b[1mProcessing frame 90 (11.88 fps)\u001b[0m\n",
      "\u001b[32m2023-12-26 21:02:26.520\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mimageflow_demo\u001b[0m:\u001b[36m41\u001b[0m - \u001b[1mProcessing frame 120 (11.82 fps)\u001b[0m\n",
      "\u001b[32m2023-12-26 21:02:26.636\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mimageflow_demo\u001b[0m:\u001b[36m120\u001b[0m - \u001b[1msave results to ./site1_1-result/track_vis/c3/2023_12_26_21_02_00.txt\u001b[0m\n",
      "\u001b[32m2023-12-26 21:02:26.720\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mimageflow_demo\u001b[0m:\u001b[36m29\u001b[0m - \u001b[1mvideo save_path is ./site1_1-result/track_vis/c1/inference_vid.mp4\u001b[0m\n",
      "\u001b[32m2023-12-26 21:02:26.724\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mimageflow_demo\u001b[0m:\u001b[36m41\u001b[0m - \u001b[1mProcessing frame 0 (100000.00 fps)\u001b[0m\n",
      "\u001b[32m2023-12-26 21:02:29.728\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mimageflow_demo\u001b[0m:\u001b[36m41\u001b[0m - \u001b[1mProcessing frame 30 (12.69 fps)\u001b[0m\n",
      "\u001b[32m2023-12-26 21:02:32.699\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mimageflow_demo\u001b[0m:\u001b[36m41\u001b[0m - \u001b[1mProcessing frame 60 (12.41 fps)\u001b[0m\n",
      "\u001b[32m2023-12-26 21:02:35.812\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mimageflow_demo\u001b[0m:\u001b[36m41\u001b[0m - \u001b[1mProcessing frame 90 (12.13 fps)\u001b[0m\n",
      "\u001b[32m2023-12-26 21:02:39.052\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mimageflow_demo\u001b[0m:\u001b[36m41\u001b[0m - \u001b[1mProcessing frame 120 (11.83 fps)\u001b[0m\n",
      "\u001b[32m2023-12-26 21:02:39.807\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36mimageflow_demo\u001b[0m:\u001b[36m120\u001b[0m - \u001b[1msave results to ./site1_1-result/track_vis/c1/2023_12_26_21_02_00.txt\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "\n",
    "\n",
    "if args.save_result:\n",
    "    output_dir = osp.join(args.save_result)\n",
    "    os.makedirs(output_dir, exist_ok=True)\n",
    "    vis_folder = osp.join(output_dir, \"track_vis\")\n",
    "    os.makedirs(vis_folder, exist_ok=True)\n",
    "\n",
    "vehicle_config = \"/data/autocon/detection/weights/yolov8x_vehicle.py\"\n",
    "vehicle_checkpoint = '/data/autocon/detection/weights/yolov8x_vehicle.pth'\n",
    "vehicle_model = init_detector(vehicle_config, vehicle_checkpoint, device='cuda:0')\n",
    "\n",
    "worker_config = \"/data/autocon/detection/weights/yolov8x_signalman.py\"\n",
    "worker_checkpoint = '/data/autocon/detection/weights/yolov8x_signalman.pth'\n",
    "worker_model = init_detector(worker_config, worker_checkpoint, device='cuda:0')\n",
    "\n",
    "pretrained_config = \"/data/autocon/detection/mmyolo/configs/yolov8/yolov8_x_mask-refine_syncbn_fast_8xb16-500e_coco.py\"\n",
    "pretrained_checkpoint = '/data/autocon/detection/weights/yolov8_x_mask-refine_syncbn_fast_8xb16-500e_coco_20230217_120411-079ca8d1.pth'\n",
    "pretrained_model = init_detector(pretrained_config, pretrained_checkpoint, device='cuda:0')\n",
    "\n",
    "det_model = (worker_model, vehicle_model, pretrained_model)\n",
    "\n",
    "clear_output(wait=False)\n",
    "\n",
    "logger.info(\"Args: {}\".format(args))\n",
    "\n",
    "tp_model = SocialImplicit(spatial_input=CFG[\"spatial_input\"],\n",
    "                          spatial_output=CFG[\"spatial_output\"],\n",
    "                          temporal_input=CFG[\"temporal_input\"],\n",
    "                          temporal_output=CFG[\"temporal_output\"],\n",
    "                          bins=CFG[\"bins\"],\n",
    "                          noise_weight=CFG[\"noise_weight\"]).cuda()\n",
    "tp_model.load_state_dict(torch.load(args.tp_weight))\n",
    "tp_model.cuda().double()\n",
    "tp_model.eval()\n",
    "\n",
    "current_time = time.localtime()\n",
    "\n",
    "if args.demo == \"video\":\n",
    "    tracker_list = imageflow_demo(det_model, vis_folder, current_time, args, tp_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "eb8120f8-1ce7-4a07-a927-2c7ec91807eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from scipy.interpolate import interp1d\n",
    "\n",
    "def inf_interpolated_coor(data):\n",
    "    \n",
    "    # Extract middle_x and right_bottom_y from tracklets\n",
    "    middle_x = {}\n",
    "    bottom_y = {}\n",
    "    \n",
    "    for frame, tracklet in data.items():\n",
    "        mid_x = tracklet[0] + tracklet[2] / 2\n",
    "        bot_y = tracklet[1] + tracklet[3]\n",
    "        middle_x[frame] = mid_x\n",
    "        bottom_y[frame] = bot_y\n",
    "    \n",
    "    # Interpolate missing frames\n",
    "    all_frames = list(range(min(middle_x.keys()), max(middle_x.keys()) + 1))\n",
    "    middle_x_interpolator = interp1d(list(middle_x.keys()), list(middle_x.values()), kind=\"linear\", fill_value=\"extrapolate\")\n",
    "    bottom_y_interpolator = interp1d(list(bottom_y.keys()), list(bottom_y.values()), kind=\"linear\", fill_value=\"extrapolate\")\n",
    "    \n",
    "    middle_x_interpolated = {frame: middle_x_interpolator(frame) for frame in all_frames}\n",
    "    bottom_y_interpolated = {frame: bottom_y_interpolator(frame) for frame in all_frames}\n",
    "    \n",
    "    # Extract representative positions for every 5 frames\n",
    "    representative_positions = []\n",
    "    \n",
    "    for frame in range(min(all_frames), max(all_frames) + 1, 5):\n",
    "        if frame in middle_x_interpolated:\n",
    "            representative_positions.append([middle_x_interpolated[frame], bottom_y_interpolated[frame]])\n",
    "\n",
    "    return representative_positions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "f621bb58-1864-4424-82d0-057d1f650486",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def tracklet_interpolation(tracklets):\n",
    "    interpolated_tracklets = {}\n",
    "    \n",
    "    for idx,tracklet in enumerate(tracklets):\n",
    "        # Interpolation\n",
    "        interpolated_tracklet = inf_interpolated_coor(tracklet.tlwh_log)\n",
    "    \n",
    "        x_coords = []\n",
    "        y_coords = []\n",
    "        for item in interpolated_tracklet:\n",
    "            c_x, b_y= item[0],item[1]\n",
    "            x_coords.append(int(c_x))\n",
    "            y_coords.append(int(b_y))\n",
    "\n",
    "        interpolated_tracklets[tracklet.track_id] = [x_coords,y_coords]\n",
    "    \n",
    "    return interpolated_tracklets\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d1218eb1-f466-48bd-9213-53a5e8f2bbd5",
   "metadata": {},
   "source": [
    "#### 1. 경로 보간 (->20개 점 확보)\n",
    "#### 2. 보간경로 점을 input(8개) 으로 주고 경로 예측(12개)\n",
    "#### 3. ADE/FDE 평가 (**실제 트래킹 된 결과값으로 평가하기)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4754ff7-17f3-458e-ba17-d5267ba2ed36",
   "metadata": {},
   "source": [
    "# 2. Trajectory Prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "1d1746a2-110b-47f0-a696-acd2921c5c1a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from model import SocialImplicit\n",
    "from CFG import CFG"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "706bda42-f106-4baf-8f35-57049c4cce41",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/javascript": [
       "IPython.notebook.clear_all_output()"
      ],
      "text/plain": [
       "<IPython.core.display.Javascript object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "tp_model = SocialImplicit(spatial_input=CFG[\"spatial_input\"],\n",
    "                          spatial_output=CFG[\"spatial_output\"],\n",
    "                          temporal_input=CFG[\"temporal_input\"],\n",
    "                          temporal_output=CFG[\"temporal_output\"],\n",
    "                          bins=CFG[\"bins\"],\n",
    "                          noise_weight=CFG[\"noise_weight\"]).cuda()\n",
    "tp_model.load_state_dict(torch.load(args.tp_weight))\n",
    "tp_model.cuda().double()\n",
    "tp_model.eval()\n",
    "\n",
    "clear_all_output()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "e82ad54b-bd2e-4eb6-991d-67eef56a3840",
   "metadata": {},
   "outputs": [],
   "source": [
    "from yolox.tracker.tp_tracker_interpolation import calculate_v_obs, seq_to_nodes, nodes_rel_to_nodes_abs, split_array\n",
    "\n",
    "# 2. 경로 예측 & 3. 평가\n",
    "def test_ade_fde(tracklet_dic):\n",
    "\n",
    "    tracklet_keys = tracklet_dic.keys()\n",
    "    tracklet_values = tracklet_dic.values()\n",
    "    # 2-1. 경로 필터링\n",
    "    filtered_track_id = []\n",
    "    filtered_traj = []\n",
    "    filtered_gt = []\n",
    "    for idx, traj in enumerate(list(tracklet_values)):\n",
    "        if np.array(traj).shape[-1] >= 20:\n",
    "            filtered_track_id.append(list(tracklet_keys)[idx])\n",
    "            filtered_traj.append(np.array(traj)[:,-20:-12])\n",
    "            filtered_gt.append(np.array(traj)[:,-12:])\n",
    "    \n",
    "    print(filtered_track_id)\n",
    "\n",
    "    obj_traj_temp = filtered_traj\n",
    "    V_y_rel_to_abs = np.array(filtered_gt).transpose(2, 0, 1) # (12, n, 2)\n",
    "    \n",
    "    # 2-2. 경로 예측\n",
    "    KSTEPS = 20\n",
    "    # create the tensor\n",
    "    obs_traj = torch.tensor(obj_traj_temp).unsqueeze(0)\n",
    "    V_obs = calculate_v_obs(obs_traj)\n",
    "    \n",
    "    obs_traj = torch.tensor(obs_traj, device='cuda:0', dtype=torch.float64)\n",
    "    V_obs = torch.tensor(V_obs, device='cuda:0', dtype=torch.float64)\n",
    "    \n",
    "    V_obs_tmp = V_obs.permute(0, 3, 1, 2)\n",
    "    V_predx = tp_model(V_obs_tmp, obs_traj, KSTEPS=KSTEPS)\n",
    "    V_x = seq_to_nodes(obs_traj.data.cpu().numpy())\n",
    "    \n",
    "    num_of_objs = V_obs.shape[-2]\n",
    "    pred_result = []\n",
    "    \n",
    "    ade_ls = {}\n",
    "    fde_ls = {}\n",
    "    for n in range(num_of_objs):\n",
    "        ade_ls[n] = []\n",
    "        fde_ls[n] = []\n",
    "    \n",
    "    for k in range(KSTEPS):\n",
    "        V_pred = V_predx[k:k + 1, ...]\n",
    "        \n",
    "        V_pred = V_pred.permute(0, 2, 3, 1)\n",
    "        V_pred = V_pred[0]\n",
    "        final_V_pred = V_pred.data.cpu().numpy()\n",
    "        V_pred_rel_to_abs = nodes_rel_to_nodes_abs(\n",
    "            final_V_pred, V_x[-1, :, :].copy())\n",
    "    \n",
    "        for n in range(num_of_objs):\n",
    "            pred = []\n",
    "            target = []\n",
    "            number_of = []\n",
    "            pred.append(V_pred_rel_to_abs[:, n:n + 1, :])\n",
    "            target.append(V_y_rel_to_abs[:, n:n + 1, :])\n",
    "            number_of.append(1)\n",
    "            \n",
    "            pred_result.append(V_pred_rel_to_abs[:, n:n + 1, :])\n",
    "        \n",
    "            ade_ls[n].append(ade(pred, target, number_of))\n",
    "            fde_ls[n].append(fde(pred, target, number_of))\n",
    "\n",
    "    ade_bigls = []\n",
    "    fde_bigls = []\n",
    "    \n",
    "    for n in range(num_of_objs):\n",
    "        ade_bigls.append(min(ade_ls[n]))\n",
    "        fde_bigls.append(min(fde_ls[n]))\n",
    "\n",
    "\n",
    "    return ade_bigls, fde_bigls\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b700fac5-bb66-4d4e-9f35-b7b6bfd07256",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "b3574182-521b-4e51-836f-02ac3b4f0bb2",
   "metadata": {},
   "source": [
    "### TP for cam#1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "35192abe-1e1f-4f48-8f9d-2de7628a2533",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 1. 경로 보간\n",
    "interpolated_tracklets1 = tracklet_interpolation(tracker_list[0].tracked_stracks) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "c54d2266-4ede-4057-bcef-156f12b0ddec",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[4, 5, 6, 1, 8, 7]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_2627622/4224312812.py:26: UserWarning: Creating a tensor from a list of numpy.ndarrays is extremely slow. Please consider converting the list to a single numpy.ndarray with numpy.array() before converting to a tensor. (Triggered internally at  /opt/conda/conda-bld/pytorch_1639180588308/work/torch/csrc/utils/tensor_new.cpp:201.)\n",
      "  obs_traj = torch.tensor(obj_traj_temp).unsqueeze(0)\n",
      "/tmp/ipykernel_2627622/4224312812.py:29: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
      "  obs_traj = torch.tensor(obs_traj, device='cuda:0', dtype=torch.float64)\n",
      "/tmp/ipykernel_2627622/4224312812.py:30: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
      "  V_obs = torch.tensor(V_obs, device='cuda:0', dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "# 3. ADE/FDE 평가 (**실제 트래킹 된 결과값으로 평가하기)\n",
    "ade_bigls, fde_bigls = test_ade_fde(interpolated_tracklets1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "27083e65-acf1-471f-94f7-2630d15b47b7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "13.226894350877002 21.75008499580892\n"
     ]
    }
   ],
   "source": [
    "ade_ = sum(ade_bigls) / len(ade_bigls)\n",
    "fde_ = sum(fde_bigls) / len(fde_bigls)\n",
    "\n",
    "print(ade_, fde_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23da7414-945f-4646-9858-ee2f26a5360b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "aff1388c-b6d9-4d0f-b6e8-2a6a815982c4",
   "metadata": {},
   "source": [
    "### TP for cam#2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "e4584906-d579-4960-b86e-2fe939689f42",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[OT_21_(1-121), OT_19_(1-121), OT_20_(1-121), OT_26_(44-121), OT_24_(21-121)]"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tracker_list[1].tracked_stracks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "4955cb43-588d-4e59-99e3-d1c55741ca44",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[21, 19, 20]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_2627622/4224312812.py:29: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
      "  obs_traj = torch.tensor(obs_traj, device='cuda:0', dtype=torch.float64)\n",
      "/tmp/ipykernel_2627622/4224312812.py:30: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
      "  V_obs = torch.tensor(V_obs, device='cuda:0', dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "interpolated_tracklets2 = tracklet_interpolation(tracker_list[1].tracked_stracks[:-1]) \n",
    "ade_bigls, fde_bigls = test_ade_fde(interpolated_tracklets2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "1b012840-ee2d-4018-be6d-1c1fe1196c37",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "8.235141204615912 13.901694283598664\n"
     ]
    }
   ],
   "source": [
    "ade_ = sum(ade_bigls) / len(ade_bigls)\n",
    "fde_ = sum(fde_bigls) / len(fde_bigls)\n",
    "\n",
    "print(ade_, fde_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "86491583-4296-4141-8d71-9db186279d0e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "8.235141204615912 13.901694283598664\n"
     ]
    }
   ],
   "source": [
    "ade_ = sum(ade_bigls) / len(ade_bigls)\n",
    "fde_ = sum(fde_bigls) / len(fde_bigls)\n",
    "\n",
    "print(ade_, fde_)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "72ba3a47-6d93-43e5-8497-99d1cc69df5f",
   "metadata": {},
   "source": [
    "### TP for cam#3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "e8dd0b67-4761-4d1b-a03a-6812c8e474e8",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[31, 32, 33, 36, 34, 35]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_2627622/4224312812.py:29: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
      "  obs_traj = torch.tensor(obs_traj, device='cuda:0', dtype=torch.float64)\n",
      "/tmp/ipykernel_2627622/4224312812.py:30: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
      "  V_obs = torch.tensor(V_obs, device='cuda:0', dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "interpolated_tracklets3 = tracklet_interpolation(tracker_list[2].tracked_stracks) \n",
    "ade_bigls, fde_bigls = test_ade_fde(interpolated_tracklets3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "6c731ccb-706a-4d1c-bf97-6b6151a5e387",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "28.46583534603135 51.08764662147622\n"
     ]
    }
   ],
   "source": [
    "ade_ = sum(ade_bigls) / len(ade_bigls)\n",
    "fde_ = sum(fde_bigls) / len(fde_bigls)\n",
    "\n",
    "print(ade_, fde_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "500eba2f-7834-4ddf-9254-aa0105b1310e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[132.34395966914516,\n",
       " 19.371333540410408,\n",
       " 3.968665106888702,\n",
       " 11.83342443565467,\n",
       " 1.137809992957485,\n",
       " 2.139819331131701]"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ade_bigls"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "2028ba16-d5db-4a90-89f8-9d56e2f27e8d",
   "metadata": {},
   "outputs": [],
   "source": [
    "del ade_bigls[0]\n",
    "del fde_bigls[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "b8fbcea2-1b72-4c5f-b5a9-587a0c055175",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[19.371333540410408,\n",
       " 3.968665106888702,\n",
       " 11.83342443565467,\n",
       " 1.137809992957485,\n",
       " 2.139819331131701]"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ade_bigls"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "1359db42-f01d-4a80-bae6-57514e2a857f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "7.690210481408593 13.549936869997193\n"
     ]
    }
   ],
   "source": [
    "ade_ = sum(ade_bigls) / len(ade_bigls)\n",
    "fde_ = sum(fde_bigls) / len(fde_bigls)\n",
    "\n",
    "print(ade_, fde_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "ced49f04-e7f7-4763-ad41-5226c4d2e190",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[19.371333540410408,\n",
       " 3.968665106888702,\n",
       " 11.83342443565467,\n",
       " 1.137809992957485,\n",
       " 2.139819331131701]"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ade_bigls"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "945b0a78-786a-4858-b604-3b89029e4587",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9406edc3-0a20-49ac-8a43-789424190aa5",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "mmyolo",
   "language": "python",
   "name": "mmyolo"
  },
  "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.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
