{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Detection - **YOLO**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "import numpy as np\n",
    "from pathlib import Path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "OpenCV: FFMPEG: tag 0x5634504d/'MP4V' is not supported with codec id 12 and format 'mp4 / MP4 (MPEG-4 Part 14)'\n",
      "OpenCV: FFMPEG: fallback to use tag 0x7634706d/'mp4v'\n"
     ]
    }
   ],
   "source": [
    "PROJECT_ROOT = Path().resolve().parent.parent\n",
    "INPUT_VIDEO_PATH = str(PROJECT_ROOT / \"data\" / \"recording_7\" / \"Recording_7.mp4\")\n",
    "OUTPUT_VIDEO_PATH = str(PROJECT_ROOT / \"data\" / \"recording_7\" / \"Recording_7_YOLO.mp4\")\n",
    "\n",
    "PATH_TO_WEIGHTS = str(PROJECT_ROOT / \"YOLO\" / \"yolov3.weights\")\n",
    "PATH_TO_CFG = str(PROJECT_ROOT / \"YOLO\" / \"yolov3.cfg\")\n",
    "\n",
    "cap = cv2.VideoCapture(INPUT_VIDEO_PATH)\n",
    "\n",
    "# Get video properties (width, height, and FPS) for saving the output video\n",
    "frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))\n",
    "frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))\n",
    "fps = cap.get(cv2.CAP_PROP_FPS)\n",
    "\n",
    "# Create VideoWriter object to save the new video\n",
    "output_video = cv2.VideoWriter(\n",
    "    OUTPUT_VIDEO_PATH, cv2.VideoWriter_fourcc(*\"MP4V\"), fps, (frame_width, frame_height)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tracking video saved successfully to: /home/davic/projects/IACV_project/bowling-analysis/data/recording_7/Recording_7_YOLO.mp4\n"
     ]
    }
   ],
   "source": [
    "# Load YOLO\n",
    "net = cv2.dnn.readNet(PATH_TO_WEIGHTS, PATH_TO_CFG)\n",
    "layer_names = net.getLayerNames()\n",
    "output_layers = [layer_names[i - 1] for i in net.getUnconnectedOutLayers()]\n",
    "\n",
    "while cap.isOpened():\n",
    "    ret, frame = cap.read()\n",
    "    if not ret:\n",
    "        break\n",
    "\n",
    "    # Prepare the frame for YOLO\n",
    "    blob = cv2.dnn.blobFromImage(\n",
    "        frame, 0.00392, (416, 416), (0, 0, 0), True, crop=False\n",
    "    )\n",
    "    net.setInput(blob)\n",
    "    outs = net.forward(output_layers)\n",
    "\n",
    "    # Post-process YOLO outputs\n",
    "    class_ids = []\n",
    "    confidences = []\n",
    "    boxes = []\n",
    "\n",
    "    for out in outs:\n",
    "        for detection in out:\n",
    "            scores = detection[5:]\n",
    "            class_id = np.argmax(scores)\n",
    "            confidence = scores[class_id]\n",
    "\n",
    "            if confidence > 0.5:  # You can adjust the threshold\n",
    "                center_x = int(detection[0] * frame.shape[1])\n",
    "                center_y = int(detection[1] * frame.shape[0])\n",
    "                w = int(detection[2] * frame.shape[1])\n",
    "                h = int(detection[3] * frame.shape[0])\n",
    "\n",
    "                x = int(center_x - w / 2)\n",
    "                y = int(center_y - h / 2)\n",
    "\n",
    "                boxes.append([x, y, w, h])\n",
    "                confidences.append(float(confidence))\n",
    "                class_ids.append(class_id)\n",
    "\n",
    "    # Use NMS (Non-Maximum Suppression) to filter overlapping boxes\n",
    "    indexes = cv2.dnn.NMSBoxes(boxes, confidences, 0.5, 0.4)\n",
    "\n",
    "    # Draw bounding boxes around the detected bowling pins\n",
    "    for i in range(len(boxes)):\n",
    "        if i in indexes:\n",
    "            x, y, w, h = boxes[i]\n",
    "            cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2)\n",
    "\n",
    "    # Write the frame with bounding boxes to the output video\n",
    "    output_video.write(frame)\n",
    "\n",
    "    cv2.imshow(\"Frame\", frame)\n",
    "    if cv2.waitKey(1) & 0xFF == ord(\"q\"):\n",
    "        break\n",
    "\n",
    "cap.release()\n",
    "output_video.release()\n",
    "\n",
    "cv2.destroyAllWindows()\n",
    "\n",
    "print(\"Tracking video saved successfully to:\", OUTPUT_VIDEO_PATH)"
   ]
  }
 ],
 "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.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}