{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "acc5f051",
   "metadata": {},
   "source": [
    "# Test optical flow for point tracking"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "7b7ed150",
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "from pathlib import Path\n",
    "import numpy as np\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ddf2d7c4",
   "metadata": {},
   "source": [
    "import video"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "1be48fef",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Opened: True, FPS: 59.94005994005994, Total Frames: 276.0\n"
     ]
    }
   ],
   "source": [
    "video_number = \"2\"\n",
    "# Define the relative path to the video file\n",
    "notebook_dir = Path().resolve()\n",
    "project_root = notebook_dir.parent.parent\n",
    "video_path = (\n",
    "    project_root\n",
    "    / \"data\"\n",
    "    / f\"recording_{video_number}\"\n",
    "    / f\"Recording_{video_number}.mp4\"\n",
    ")\n",
    "video_path = str(video_path)\n",
    "\n",
    "# Load the video\n",
    "cap = cv2.VideoCapture(video_path)\n",
    "\n",
    "# Check\n",
    "print(\n",
    "    f\"Opened: {cap.isOpened()}, FPS: {cap.get(cv2.CAP_PROP_FPS)}, Total Frames: {cap.get(cv2.CAP_PROP_FRAME_COUNT)}\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5289312",
   "metadata": {},
   "source": [
    "Import points"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "bb4506b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the path to the CSV file\n",
    "input_data_path = (\n",
    "    project_root\n",
    "    / \"data\"\n",
    "    / \"auxiliary_data\"\n",
    "    / \"lane_points\"\n",
    "    / f\"lane_points_processed_{video_number}.csv\"\n",
    ")\n",
    "\n",
    "# Load the CSV file into a DataFrame\n",
    "points_df = pd.read_csv(input_data_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28fa5cf9",
   "metadata": {},
   "source": [
    "Tracking of the top right point"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "b8cc04b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Reset the video to the beginning\n",
    "cap.set(cv2.CAP_PROP_POS_FRAMES, 0)\n",
    "\n",
    "# Read the first frame\n",
    "ret, old_frame = cap.read()\n",
    "if not ret:\n",
    "    print(\"Cannot read video\")\n",
    "    cap.release()\n",
    "    exit()\n",
    "\n",
    "# Convert to grayscale\n",
    "old_gray = cv2.cvtColor(old_frame, cv2.COLOR_BGR2GRAY)\n",
    "\n",
    "# Select the top-right point of the first frame\n",
    "point = (int(points_df.iloc[0][\"up_left_x\"]), int(points_df.iloc[0][\"up_left_y\"]))\n",
    "\n",
    "point_frame = old_frame.copy()\n",
    "cv2.circle(point_frame, point, 5, (0, 0, 255), 2)\n",
    "cv2.imshow(\"Frame\", point_frame)\n",
    "\n",
    "# create the array of the point to track\n",
    "old_points = np.array([[point[0], point[1]]], dtype=np.float32).reshape(-1, 1, 2)\n",
    "\n",
    "# store points\n",
    "points = [point]\n",
    "\n",
    "# Tracking loop\n",
    "while True:\n",
    "    ret, new_frame = cap.read()\n",
    "    if not ret:\n",
    "        break\n",
    "\n",
    "    new_gray = cv2.cvtColor(new_frame, cv2.COLOR_BGR2GRAY)\n",
    "\n",
    "    # Calculate optical flow\n",
    "    new_points, status, error = cv2.calcOpticalFlowPyrLK(\n",
    "        old_gray, new_gray, old_points, None\n",
    "    )\n",
    "\n",
    "    if status[0] == 1:\n",
    "        x, y = new_points.ravel()\n",
    "\n",
    "        # Update the previous frame and points\n",
    "        old_gray = new_gray.copy()\n",
    "        old_points = new_points.reshape(-1, 1, 2)\n",
    "        points.append((int(x), int(y)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e29fd2f0",
   "metadata": {},
   "source": [
    "Generate video"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "ad2b1710",
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_point_to_frame(video_frame, point, color=(0, 255, 0), radius=5):\n",
    "    \"\"\"Add a point to the video frame.\"\"\"\n",
    "    cv2.circle(video_frame, point, radius, color, -1)\n",
    "    return video_frame"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "9698c7fa",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Adjusted video saved to C:\\Users\\miche\\OneDrive\\Documenti\\GitHub\\bowling-analysis\\data\\recording_2\\Point_Tracking.mp4\n"
     ]
    }
   ],
   "source": [
    "# Reset the video to the beginning\n",
    "cap.set(cv2.CAP_PROP_POS_FRAMES, 0)\n",
    "\n",
    "# Define the codec and create a VideoWriter object to save the modified frames\n",
    "output_path = project_root / \"data\" / f\"recording_{video_number}\" / \"Point_Tracking.mp4\"\n",
    "fourcc = cv2.VideoWriter_fourcc(*\"mp4v\")  # Use 'mp4v' codec for MP4 format\n",
    "fps = int(cap.get(cv2.CAP_PROP_FPS))\n",
    "frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))\n",
    "frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))\n",
    "out = cv2.VideoWriter(str(output_path), fourcc, fps, (frame_width, frame_height))\n",
    "\n",
    "\n",
    "# Loop through each frame in the video\n",
    "frame_index = 0\n",
    "while frame_index < len(points):\n",
    "    ret, video_frame = cap.read()\n",
    "    if not ret:\n",
    "        print(\"End of video or failed to read the frame at iteration\", frame_index)\n",
    "        break\n",
    "\n",
    "    # draw the lines on the frame\n",
    "    modified_frame = add_point_to_frame(\n",
    "        video_frame, points[frame_index], color=(0, 255, 0), radius=5\n",
    "    )\n",
    "\n",
    "    # Write the modified frame to the output video\n",
    "    out.write(modified_frame)\n",
    "\n",
    "    # Increment the frame index\n",
    "    frame_index += 1\n",
    "\n",
    "# Release the video capture and writer objects\n",
    "# cap.release()\n",
    "out.release()\n",
    "\n",
    "print(f\"Adjusted video saved to {output_path}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}