{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a4b9809a",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/anaconda3/envs/yt38/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"
     ]
    },
    {
     "ename": "ModuleNotFoundError",
     "evalue": "No module named 'ensemble_boxes'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[1], line 11\u001b[0m\n\u001b[1;32m      9\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mtime\u001b[39;00m\n\u001b[1;32m     10\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mPIL\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Image\n\u001b[0;32m---> 11\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mensemble_boxes\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;241m*\u001b[39m\n",
      "\u001b[0;31mModuleNotFoundError\u001b[0m: No module named 'ensemble_boxes'"
     ]
    }
   ],
   "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",
    "\n",
    "import pandas as pd\n",
    "import time\n",
    "from PIL import Image\n",
    "from ensemble_boxes import *"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "91e6e452",
   "metadata": {},
   "source": [
    "# Test "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "efd6f33b",
   "metadata": {},
   "outputs": [],
   "source": [
    "cctv_path = './video/cctv/'\n",
    "youtube_path = './video/youtube/'\n",
    "\n",
    "cctv_vid_list = [file for file in os.listdir(cctv_path) if file.endswith(\".mp4\") or file.endswith(\".avi\")]\n",
    "youtube_vid_list = [file for file in os.listdir(youtube_path) if file.endswith(\".mp4\") or file.endswith(\".avi\")]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da1d0e05",
   "metadata": {},
   "outputs": [],
   "source": [
    "model1 = torch.hub.load(\"/data/kwater/yolov5\",\n",
    "                       'custom', path='/data/kwater/test/weight/5s_0.68.pt',\n",
    "                       source='local',device=\"cuda:0\")  # local repo\n",
    "\n",
    "model2 = torch.hub.load(\"/data/kwater/yolov5\",\n",
    "                       'custom', path='/data/kwater/test/weight/5m_0.80.pt',\n",
    "                       source='local',device=\"cuda:0\")  # local repo\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f33d8ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "def draw_bbox(img, vehicle_results, thres):\n",
    "    for obj in vehicle_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",
    "            color = (255,0,0)\n",
    "\n",
    "            tk = 1\n",
    "            # if obj['risk1'] == 'danger':\n",
    "            #     tk = 5\n",
    "            \n",
    "            class_dic = {0:\"Leak\",\n",
    "                        1:\"Pole\"}\n",
    "            content =  class_dic[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": null,
   "id": "74f68d1d",
   "metadata": {},
   "outputs": [],
   "source": [
    "cctv_vid_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55c4fd9e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def convert_wbf(results, threshold, image_size):\n",
    "    boxes_list = []\n",
    "    scores_list = []\n",
    "    labels_list = []\n",
    "    \n",
    "    for preds in results:\n",
    "        if preds[-2] >= threshold:\n",
    "            box = preds[:4]\n",
    "            boxes_list.append([box[0]/image_size[1],\n",
    "                               box[1]/image_size[0],\n",
    "                               box[2]/image_size[1],\n",
    "                               box[3]/image_size[0]]\n",
    "                               )\n",
    "            score = preds[-2].astype(float)\n",
    "            scores_list.append(score)\n",
    "            labels_list.append(preds[-1])\n",
    "    return boxes_list, labels_list, scores_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b8f0491",
   "metadata": {},
   "outputs": [],
   "source": [
    "frame_results = []\n",
    "for file in cctv_vid_list:\n",
    "#     print(file)\n",
    "#     print(file.split('.')[0]+'.csv')\n",
    "    \n",
    "    csv_file = os.path.join(cctv_path,file.split('.')[0]+'.csv')\n",
    "    data = np.array(pd.read_csv(csv_file, header=None))\n",
    "    \n",
    "    video_file = os.path.join(cctv_path,file)\n",
    "    cap = cv2.VideoCapture(video_file) \n",
    "    ret, frame = cap.read()   \n",
    "    \n",
    "    video_length = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))\n",
    "    video_fps = int(cap.get(cv2.CAP_PROP_FPS))\n",
    "    image_size = frame.shape\n",
    "    \n",
    "    # image array\n",
    "    image_array = []\n",
    "    count = 0\n",
    "\n",
    "    while frame is not None:\n",
    "\n",
    "#         frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)\n",
    "\n",
    "        results1 = model1(frame)\n",
    "        bbox_results1=results1.xyxy[0].cpu().numpy() ## 결과값\n",
    "        \n",
    "        results2 = model2(frame)\n",
    "        bbox_results2=results2.xyxy[0].cpu().numpy() ## 결과값\n",
    "        \n",
    "        ## input format is xyxy\n",
    "        threshold = 0.4\n",
    "        boxes_list1, labels_list1, scores_list1 = convert_wbf(bbox_results1, threshold = threshold, image_size = image_size)\n",
    "        boxes_list2, labels_list2, scores_list2 = convert_wbf(bbox_results2, threshold = threshold, image_size = image_size)\n",
    "        \n",
    "        boxes_list_merges = []\n",
    "        labels_list_merges = []\n",
    "        scores_list_merges = []\n",
    "        boxes_list_merges.append(boxes_list1)\n",
    "        boxes_list_merges.append(boxes_list2)\n",
    "        labels_list_merges.append(labels_list1)\n",
    "        labels_list_merges.append(labels_list2)\n",
    "        scores_list_merges.append(scores_list1)\n",
    "        scores_list_merges.append(scores_list2)\n",
    "        \n",
    "        weights = [1, 1]\n",
    "        iou_thr = 0.5\n",
    "        skip_box_thr = 0.0001\n",
    "        sigma = 0.1\n",
    "        boxes, scores, labels = weighted_boxes_fusion(boxes_list_merges, scores_list_merges, \n",
    "                                                      labels_list_merges, weights=weights, iou_thr=iou_thr, skip_box_thr=skip_box_thr)\n",
    "        \n",
    "        img_w = image_size[1]\n",
    "        img_h = image_size[0]\n",
    "        \n",
    "        ## float to bbox\n",
    "        for idx,row in enumerate(boxes):\n",
    "            boxes[idx][0] *= img_w\n",
    "            boxes[idx][1] *= img_h\n",
    "            boxes[idx][2] *= img_w\n",
    "            boxes[idx][3] *= img_h\n",
    "            \n",
    "        ensemble_bbox_results = []\n",
    "        for idx,box in enumerate(boxes):\n",
    "            temp=[]\n",
    "            temp.extend(box)\n",
    "            temp.extend([scores[idx]])\n",
    "            temp.extend([int(labels[idx])])\n",
    "            \n",
    "            ensemble_bbox_results.append(temp)\n",
    "        ensemble_bbox_results = np.array(ensemble_bbox_results)\n",
    "            \n",
    "#         result_img=draw_bbox(frame.copy(), ensemble_bbox_results, 0.25)\n",
    "        \n",
    "        ######### SCORING\n",
    "        # COMPARE AND SCORING\n",
    "        gt = []\n",
    "        for row in data:\n",
    "            if row[0] == count:\n",
    "                gt.append(list(row[1:]))\n",
    "                \n",
    "        gt = np.array(gt)\n",
    "        \n",
    "        # bbox to float\n",
    "        if len(ensemble_bbox_results)>0:\n",
    "            pred_bb1 = np.array(ensemble_bbox_results[:,:4], dtype=np.float64)\n",
    "            pred_cls1 = np.array(ensemble_bbox_results[:,5], dtype=np.int64)\n",
    "            pred_conf1 = np.array(np.round(ensemble_bbox_results[:,4],2), dtype=np.float64)\n",
    "            for idx,row in enumerate(pred_bb1):\n",
    "                pred_bb1[idx][0] /= img_w\n",
    "                pred_bb1[idx][1] /= img_h\n",
    "                pred_bb1[idx][2] /= img_w\n",
    "                pred_bb1[idx][3] /= img_h\n",
    "        else:\n",
    "            pred_bb1 = np.array([])\n",
    "            pred_cls1 = np.array([])\n",
    "            pred_conf1 = np.array([])\n",
    "            \n",
    "        if len(gt)>0:\n",
    "            gt_bb1 = np.array(gt[:,:4], dtype=np.float64)\n",
    "            gt_cls1 = np.array(gt[:,4], dtype=np.int64)\n",
    "            for idx,row in enumerate(gt_bb1):\n",
    "                gt_bb1[idx][0] /= img_w\n",
    "                gt_bb1[idx][1] /= img_h\n",
    "                gt_bb1[idx][2] /= img_w\n",
    "                gt_bb1[idx][3] /= img_h  \n",
    "        else:\n",
    "            gt_bb1 = np.array([])\n",
    "            gt_cls1 = np.array([])\n",
    "            \n",
    "        frame_results.append([pred_bb1, pred_cls1, pred_conf1, gt_bb1, gt_cls1])\n",
    "\n",
    "        clear_output(wait=True)\n",
    "#         %matplotlib inline\n",
    "#         plt.imshow(result_img)\n",
    "#         plt.show()\n",
    "        \n",
    "#         print(results)\n",
    "#         print(bbox_results)\n",
    "#         print(np.array(data)[0][1:])\n",
    "\n",
    "#         image_array.append(result_img)\n",
    "\n",
    "        count+=1\n",
    "        ret, frame = cap.read()\n",
    "\n",
    "        print(\"total L : \", video_length)\n",
    "        print(\"now : \",count)\n",
    "        \n",
    "        \n",
    "#     out = cv2.VideoWriter(\"./cctv_results/inferenced_\"+file.split('.')[0]+\".mp4\",\n",
    "#                       cv2.VideoWriter_fourcc(*'DIVX'), video_fps, (image_size[1],image_size[0]))\n",
    "\n",
    "#     for i in range(len(image_array)):\n",
    "#         out.write(image_array[i])\n",
    "#     out.release()\n",
    "    \n",
    "    \n",
    "    \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eba59b1b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from mean_average_precision.detection_map import DetectionMAP\n",
    "from mean_average_precision.utils.show_frame import show_frame\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6bb3610a",
   "metadata": {},
   "outputs": [],
   "source": [
    "n_class = 2\n",
    "\n",
    "mAP = DetectionMAP(n_class)\n",
    "for i, frame in enumerate(frame_results):\n",
    "#     print(\"Evaluate frame {}\".format(i))\n",
    "#     show_frame(*frame)\n",
    "    mAP.evaluate(*frame)\n",
    "\n",
    "# class_index\n",
    "# %matplotlib inline\n",
    "mAP.plot()\n",
    "# plt.show()\n",
    "plt.savefig(\"pr_curve_example.png\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf2dbe3a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# pred : [[     612.35      267.11        1017       567.1      0.9717     0]]\n",
    "# gt : [[  609  269 1021  571    0]]\n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3fc76880",
   "metadata": {},
   "outputs": [],
   "source": [
    "frame_results = []\n",
    "for file in youtube_vid_list:\n",
    "#     print(file)\n",
    "#     print(file.split('.')[0]+'.csv')\n",
    "    \n",
    "    csv_file = os.path.join(youtube_path,file.split('.')[0]+'.csv')\n",
    "    data = np.array(pd.read_csv(csv_file, header=None))\n",
    "    \n",
    "    video_file = os.path.join(youtube_path,file)\n",
    "    cap = cv2.VideoCapture(video_file) \n",
    "    ret, frame = cap.read()   \n",
    "    \n",
    "    video_length = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))\n",
    "    video_fps = int(cap.get(cv2.CAP_PROP_FPS))\n",
    "    image_size = frame.shape\n",
    "    \n",
    "    # image array\n",
    "    image_array = []\n",
    "    count = 0\n",
    "\n",
    "    while frame is not None:\n",
    "\n",
    "#         frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)\n",
    "\n",
    "        results1 = model1(frame)\n",
    "        bbox_results1=results1.xyxy[0].cpu().numpy() ## 결과값\n",
    "        \n",
    "        results2 = model2(frame)\n",
    "        bbox_results2=results2.xyxy[0].cpu().numpy() ## 결과값\n",
    "        \n",
    "        ## input format is xyxy\n",
    "        threshold = 0.4\n",
    "        boxes_list1, labels_list1, scores_list1 = convert_wbf(bbox_results1, threshold = threshold, image_size = image_size)\n",
    "        boxes_list2, labels_list2, scores_list2 = convert_wbf(bbox_results2, threshold = threshold, image_size = image_size)\n",
    "        \n",
    "        boxes_list_merges = []\n",
    "        labels_list_merges = []\n",
    "        scores_list_merges = []\n",
    "        boxes_list_merges.append(boxes_list1)\n",
    "        boxes_list_merges.append(boxes_list2)\n",
    "        labels_list_merges.append(labels_list1)\n",
    "        labels_list_merges.append(labels_list2)\n",
    "        scores_list_merges.append(scores_list1)\n",
    "        scores_list_merges.append(scores_list2)\n",
    "        \n",
    "        weights = [1, 1]\n",
    "        iou_thr = 0.5\n",
    "        skip_box_thr = 0.0001\n",
    "        sigma = 0.1\n",
    "        boxes, scores, labels = weighted_boxes_fusion(boxes_list_merges, scores_list_merges, \n",
    "                                                      labels_list_merges, weights=weights, iou_thr=iou_thr, skip_box_thr=skip_box_thr)\n",
    "        \n",
    "        img_w = image_size[1]\n",
    "        img_h = image_size[0]\n",
    "        \n",
    "        ## float to bbox\n",
    "        for idx,row in enumerate(boxes):\n",
    "            boxes[idx][0] *= img_w\n",
    "            boxes[idx][1] *= img_h\n",
    "            boxes[idx][2] *= img_w\n",
    "            boxes[idx][3] *= img_h\n",
    "            \n",
    "        ensemble_bbox_results = []\n",
    "        for idx,box in enumerate(boxes):\n",
    "            temp=[]\n",
    "            temp.extend(box)\n",
    "            temp.extend([scores[idx]])\n",
    "            temp.extend([int(labels[idx])])\n",
    "            \n",
    "            ensemble_bbox_results.append(temp)\n",
    "        ensemble_bbox_results = np.array(ensemble_bbox_results)\n",
    "            \n",
    "#         result_img=draw_bbox(frame.copy(), ensemble_bbox_results, 0.25)\n",
    "        \n",
    "        ######### SCORING\n",
    "        # COMPARE AND SCORING\n",
    "        gt = []\n",
    "        for row in data:\n",
    "            if row[0] == count:\n",
    "                gt.append(list(row[1:]))\n",
    "                \n",
    "        gt = np.array(gt)\n",
    "        \n",
    "        # bbox to float\n",
    "        if len(ensemble_bbox_results)>0:\n",
    "            pred_bb1 = np.array(ensemble_bbox_results[:,:4], dtype=np.float64)\n",
    "            pred_cls1 = np.array(ensemble_bbox_results[:,5], dtype=np.int64)\n",
    "            pred_conf1 = np.array(np.round(ensemble_bbox_results[:,4],2), dtype=np.float64)\n",
    "            for idx,row in enumerate(pred_bb1):\n",
    "                pred_bb1[idx][0] /= img_w\n",
    "                pred_bb1[idx][1] /= img_h\n",
    "                pred_bb1[idx][2] /= img_w\n",
    "                pred_bb1[idx][3] /= img_h\n",
    "        else:\n",
    "            pred_bb1 = np.array([])\n",
    "            pred_cls1 = np.array([])\n",
    "            pred_conf1 = np.array([])\n",
    "            \n",
    "        if len(gt)>0:\n",
    "            gt_bb1 = np.array(gt[:,:4], dtype=np.float64)\n",
    "            gt_cls1 = np.array(gt[:,4], dtype=np.int64)\n",
    "            for idx,row in enumerate(gt_bb1):\n",
    "                gt_bb1[idx][0] /= img_w\n",
    "                gt_bb1[idx][1] /= img_h\n",
    "                gt_bb1[idx][2] /= img_w\n",
    "                gt_bb1[idx][3] /= img_h  \n",
    "        else:\n",
    "            gt_bb1 = np.array([])\n",
    "            gt_cls1 = np.array([])\n",
    "            \n",
    "        frame_results.append([pred_bb1, pred_cls1, pred_conf1, gt_bb1, gt_cls1])\n",
    "\n",
    "        clear_output(wait=True)\n",
    "#         %matplotlib inline\n",
    "#         plt.imshow(result_img)\n",
    "#         plt.show()\n",
    "        \n",
    "#         print(results)\n",
    "#         print(bbox_results)\n",
    "#         print(np.array(data)[0][1:])\n",
    "\n",
    "#         image_array.append(result_img)\n",
    "\n",
    "        count+=1\n",
    "        ret, frame = cap.read()\n",
    "\n",
    "        print(\"total L : \", video_length)\n",
    "        print(\"now : \",count)\n",
    "        \n",
    "        \n",
    "#     out = cv2.VideoWriter(\"./cctv_results/inferenced_\"+file.split('.')[0]+\".mp4\",\n",
    "#                       cv2.VideoWriter_fourcc(*'DIVX'), video_fps, (image_size[1],image_size[0]))\n",
    "\n",
    "#     for i in range(len(image_array)):\n",
    "#         out.write(image_array[i])\n",
    "#     out.release()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79a516b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from mean_average_precision.detection_map import DetectionMAP\n",
    "from mean_average_precision.utils.show_frame import show_frame\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a429143d",
   "metadata": {},
   "outputs": [],
   "source": [
    "n_class = 2\n",
    "\n",
    "mAP = DetectionMAP(n_class)\n",
    "for i, frame in enumerate(frame_results):\n",
    "#     print(\"Evaluate frame {}\".format(i))\n",
    "#     show_frame(*frame)\n",
    "    mAP.evaluate(*frame)\n",
    "\n",
    "# class_index\n",
    "mAP.plot()\n",
    "plt.show()\n",
    "plt.savefig(\"pr_curve_example.png\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84681ff5",
   "metadata": {},
   "outputs": [],
   "source": [
    "temp = [[[]]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cee6cbdf-3104-4971-8273-84675535f16b",
   "metadata": {},
   "outputs": [],
   "source": [
    "temp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b431991e-db76-4d8e-847d-f882e0a1857a",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.array(temp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea3ab8eb-10a7-4c97-a03b-55bb7eacd3b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "list(a)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b905029f-306a-4eaa-b5b3-fffe7eac95ff",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "yt38",
   "language": "python",
   "name": "yt38"
  },
  "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
}