tracking_vscode.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from ultralytics import YOLO\n",
    "import json\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# SETUP (Запускать только после клонирования репозитория)\n",
    "\n",
    "os.mkdir(\"jsons\")\n",
    "os.mkdir(\"result\")\n",
    "os.mkdir(\"vid\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Empty DataFrame\n",
      "Columns: [file_name, car, quantity_car, average_speed_car, van, quantity_van, average_speed_van, bus, quantity_bus, average_speed_bus]\n",
      "Index: []\n",
      "Доступные видео:  ['KRA-21-62-2023-08-07-morning.mp4', 'KRA-25-73-2023-08-24-morning.mp4', 'KRA-43-164-2023-08-08-evening.mp4']\n",
      "=== KRA-21-62-2023-08-07-morning.mp4 - 29.0 FPS ===\n",
      "8 = 8 - 0\n",
      "Средняя скорость по классу: {2: 58.0}\n",
      "Финальное количество объектов в видео KRA-21-62-2023-08-07-morning.mp4: {2: 1, 5: 0, 7: 0}\n",
      "=== KRA-25-73-2023-08-24-morning.mp4 - 30.0 FPS ===\n",
      "5 = 5 - 0\n",
      "Средняя скорость по классу: {7: 90.0}\n",
      "Финальное количество объектов в видео KRA-25-73-2023-08-24-morning.mp4: {2: 0, 5: 0, 7: 1}\n",
      "=== KRA-43-164-2023-08-08-evening.mp4 - 30.0 FPS ===\n",
      "5 = 5 - 0\n",
      "Средняя скорость по классу: {7: 90.0}\n",
      "Финальное количество объектов в видео KRA-43-164-2023-08-08-evening.mp4: {2: 0, 5: 0, 7: 1}\n",
      "                       file_name  car  quantity_car  average_speed_car  van  \\\n",
      "0   KRA-21-62-2023-08-07-morning  car             1               58.0  van   \n",
      "1   KRA-25-73-2023-08-24-morning  car             0                0.0  van   \n",
      "2  KRA-43-164-2023-08-08-evening  car             0                0.0  van   \n",
      "\n",
      "   quantity_van  average_speed_van  bus  quantity_bus  average_speed_bus  \n",
      "0             0                0.0  bus             0                  0  \n",
      "1             1               90.0  bus             0                  0  \n",
      "2             1               90.0  bus             0                  0  \n"
     ]
    }
   ],
   "source": [
    "class Transport:\n",
    "    def __init__(self, idx: str) -> None:\n",
    "        self.first_line: bool = False\n",
    "        self.second_line: bool = False\n",
    "        self.inside_area: bool = False\n",
    "\n",
    "        self.idx: str = idx\n",
    "        self.ts_type: list = []\n",
    "\n",
    "        self.frames: list = []\n",
    "\n",
    "\n",
    "def get_crop_x(annotated_frame, data):\n",
    "    areas = data[\"areas\"][0]\n",
    "    x_min = 10000\n",
    "    x_max = 0\n",
    "    for point in areas:\n",
    "        x, y = int(point[0] * annotated_frame.shape[1]), int(point[1] * annotated_frame.shape[0])\n",
    "        if x < x_min:\n",
    "            x_min = x\n",
    "        if x > x_max:\n",
    "            x_max = x\n",
    "\n",
    "    return x_min, x_max\n",
    "\n",
    "\n",
    "def is_inside_zone(point_x, point_y, frame) -> bool:\n",
    "    areas = []\n",
    "    areas1 = data['areas'][0]\n",
    "\n",
    "    p1 = int(areas1[0][0] * frame.shape[1]), int(areas1[0][1] * frame.shape[0])\n",
    "    p2 = int(areas1[1][0] * frame.shape[1]), int(areas1[1][1] * frame.shape[0])\n",
    "    p3 = int(areas1[2][0] * frame.shape[1]), int(areas1[2][1] * frame.shape[0])\n",
    "    p4 = int(areas1[3][0] * frame.shape[1]), int(areas1[3][1] * frame.shape[0])\n",
    "\n",
    "    areas.append(np.array([p1, p2, p3, p4]))\n",
    "\n",
    "    try:\n",
    "        areas2 = data['areas'][1]\n",
    "\n",
    "        p1 = int(areas2[0][0] * frame.shape[1]), int(areas2[0][1] * frame.shape[0])\n",
    "        p2 = int(areas2[1][0] * frame.shape[1]), int(areas2[1][1] * frame.shape[0])\n",
    "        p3 = int(areas2[2][0] * frame.shape[1]), int(areas2[2][1] * frame.shape[0])\n",
    "        p4 = int(areas2[3][0] * frame.shape[1]), int(areas2[3][1] * frame.shape[0])\n",
    "        areas.append(np.array([p1, p2, p3, p4]))\n",
    "    except:\n",
    "        pass\n",
    "\n",
    "    for area in areas:\n",
    "        polyTestRes = cv2.pointPolygonTest(area, (point_x, point_y), measureDist=False)\n",
    "        cv2.polylines(frame, pts=[area], isClosed=True, color=(0, 0, 255))\n",
    "        if polyTestRes >= 0:\n",
    "            return True\n",
    "    return False\n",
    "\n",
    "\n",
    "FRAME_ANALYSIS_FREQ = 4\n",
    "VIDEO_DIR = \"vid\"\n",
    "JSON_DIR = \"jsons\"\n",
    "USED_CLASSES = [2, 5, 7]  # 2 - car, 5 - bus, 7 - truck\n",
    "\n",
    "init_keys = [\"file_name\", \"car\", \"quantity_car\", \"average_speed_car\", \"van\", \"quantity_van\", \"average_speed_van\", \"bus\",\n",
    "             \"quantity_bus\", \"average_speed_bus\"]\n",
    "result_df = pd.DataFrame(dict.fromkeys(init_keys, []))\n",
    "print(result_df)\n",
    "\n",
    "vids = os.listdir(VIDEO_DIR)\n",
    "print(\"Доступные видео: \", vids)\n",
    "for vid_fname_idx in range(len(vids)):\n",
    "    vid_fname = vids[vid_fname_idx]\n",
    "    json_name = vid_fname.split('.mp4')[0]\n",
    "\n",
    "    with open(f'{JSON_DIR}/{json_name}.json', 'r', encoding='utf-8') as f:\n",
    "        data = json.load(f)\n",
    "\n",
    "    model = YOLO('yolov8x.pt')\n",
    "\n",
    "    vid_path = f\"{VIDEO_DIR}/{vid_fname}\"\n",
    "    vid_capture = cv2.VideoCapture(vid_path)\n",
    "    vid_fps = vid_capture.get(cv2.CAP_PROP_FPS)\n",
    "    print(f\"=== {vid_fname} - {vid_fps} FPS ===\")\n",
    "\n",
    "    appeared_ids = []\n",
    "\n",
    "    transport_dict: dict = {}\n",
    "\n",
    "    frame_idx = 0\n",
    "    frame_idx_in_cond = 0\n",
    "\n",
    "    while vid_capture.isOpened():\n",
    "        success, frame = vid_capture.read()\n",
    "\n",
    "        if success:\n",
    "            if frame_idx % FRAME_ANALYSIS_FREQ == 0:\n",
    "                frame = cv2.resize(frame, (640, 360)) # (1280, 720)\n",
    "\n",
    "                width = frame.shape[1]\n",
    "                x_min, x_max = get_crop_x(annotated_frame=frame, data=data)\n",
    "\n",
    "                new_width = width - (x_min + x_max)\n",
    "                start_x = x_min\n",
    "                cropped_frame = frame[:, start_x:start_x + new_width]\n",
    "                color = (0, 0, 255)\n",
    "\n",
    "                # !!! Изменить verbose на True если нужна инфа от модели !!!\n",
    "                track_results = model.track(frame,\n",
    "                                            persist=True,\n",
    "                                            classes=USED_CLASSES,\n",
    "                                            verbose=False,\n",
    "                                            conf=0.5,\n",
    "                                            imgsz=(384, 640)) \n",
    "                annotated_frame = track_results[0].plot()\n",
    "                try:\n",
    "                    bb_on_frame_ids: np.ndarray = track_results[0].boxes.id.numpy()\n",
    "                    bb_on_frame_cls: np.ndarray = track_results[0].boxes.cls.numpy()\n",
    "                    bb_center: np.ndarray = track_results[0].boxes.xywh.numpy()\n",
    "                    bb_corners: np.ndarray = track_results[0].boxes.xyxy.numpy()\n",
    "\n",
    "                    for i in range(len(bb_on_frame_ids)):\n",
    "\n",
    "                        idx = str(int(bb_on_frame_ids[i]))\n",
    "\n",
    "                        if idx not in transport_dict.keys():\n",
    "                            transport_dict[idx] = Transport(str(idx))\n",
    "\n",
    "                        cv2.circle(annotated_frame, (int(bb_center[i][0]), int(bb_center[i][1] + (bb_center[i][3] / 2))), 8, (232, 88, 163), -1)\n",
    "\n",
    "                        if is_inside_zone(bb_center[i][0], bb_center[i][1] + (bb_center[i][3] / 2), annotated_frame):\n",
    "                            transport_dict[idx].frames.append(frame_idx_in_cond)\n",
    "                            transport_dict[idx].ts_type.append(int(bb_on_frame_cls[i]))\n",
    "                            transport_dict[idx].inside_area = True\n",
    "\n",
    "                except Exception as e:\n",
    "                    print(e)\n",
    "\n",
    "                annotated_frame = cv2.resize(annotated_frame, (1280, 640))\n",
    "                cv2.imshow(\"YOLOv8 Tracking\", annotated_frame)\n",
    "                key = cv2.waitKey(1)\n",
    "                frame_idx_in_cond += 1\n",
    "\n",
    "                if key == 27:\n",
    "                    break\n",
    "            frame_idx += 1\n",
    "        else:\n",
    "            break\n",
    "\n",
    "    class_average_speed: dict = {}\n",
    "    count_dict = dict.fromkeys(USED_CLASSES, 0)\n",
    "\n",
    "    for k in transport_dict.keys():\n",
    "        ts: Transport = transport_dict[k]\n",
    "\n",
    "        if len(ts.frames) < 6:\n",
    "            continue\n",
    "\n",
    "        import operator\n",
    "\n",
    "        d = dict.fromkeys(ts.ts_type)\n",
    "        for k in d.keys():\n",
    "            d[k] = ts.ts_type.count(k)\n",
    "        ts_type = max(d.items(), key=operator.itemgetter(1))[0]\n",
    "\n",
    "        if ts.inside_area:\n",
    "            count_dict[ts_type] += 1\n",
    "\n",
    "        if ts.inside_area:\n",
    "            delta_f = ts.frames[-1] - ts.frames[0]\n",
    "            print(f'{delta_f} = {ts.frames[-1]} - {ts.frames[0]}')\n",
    "            m: float = 20\n",
    "            ms_to_kmh: float = 3.6\n",
    "            k: float = (vid_fps / FRAME_ANALYSIS_FREQ)\n",
    "            average_speed = m / ((delta_f + 1) / k) * ms_to_kmh\n",
    "            class_average_speed.setdefault(ts_type, []).append(average_speed)\n",
    "\n",
    "    for k in class_average_speed.keys():\n",
    "        class_average_speed[k] = sum(class_average_speed[k]) / len(class_average_speed[k])\n",
    "\n",
    "    print(f\"Средняя скорость по классу: {class_average_speed}\")\n",
    "    print(f\"Финальное количество объектов в видео {vid_fname}: {count_dict}\")\n",
    "    vid_fname_formatless = vid_fname.split(\".\")[0]\n",
    "\n",
    "    try:\n",
    "        car_avg_spd = class_average_speed[2]\n",
    "    except KeyError: car_avg_spd = 0\n",
    "    try:\n",
    "        bus_avg_spd = class_average_speed[5]\n",
    "    except KeyError: bus_avg_spd = 0\n",
    "    try:\n",
    "        van_avg_spd = class_average_speed[7]\n",
    "    except KeyError: van_avg_spd = 0\n",
    "\n",
    "    data_inserted = [vid_fname_formatless, \"car\", count_dict[2], car_avg_spd, \"van\", count_dict[7], van_avg_spd, \"bus\", count_dict[5], bus_avg_spd]\n",
    "\n",
    "    result_df.loc[vid_fname_idx] = data_inserted\n",
    "    result_df.to_excel('./result/submission.xlsx', index=False)\n",
    "    result_df.to_csv('./result/submission.csv', index=False)\n",
    "\n",
    "    vid_capture.release()\n",
    "    cv2.destroyAllWindows()\n",
    "\n",
    "print(result_df)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "language": "python",
   "name": "python3"
  },
  "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.11.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}