main.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# Vision\n",
    "\n",
    "The purpose of this notebook is to create images that will make the workflow of moving apples easy to do. \n",
    "The final product will be a Streamlit site where I can go forward and backwards and see frames and the transitions between them.\n",
    "\n",
    "Requirements: Original frame from the video, the frame downsampled to circles with coordinates of every apple, and a diff frame with the added apples in green and removed apples in red.\n",
    "\n",
    "There was some early experimentation with the travelling salesman problem"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "6747e560b2b692fd"
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "outputs": [],
   "source": [
    "import cv2\n",
    "import numpy as np\n",
    "from pathlib import Path\n",
    "import os"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-02-24T18:40:43.151640700Z",
     "start_time": "2024-02-24T18:40:43.148640Z"
    }
   },
   "id": "23178aabc4fd075"
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "width: 960, height: 720, fps: 30, n_frames: 6572\n"
     ]
    }
   ],
   "source": [
    "badapple_path = Path(\".\").resolve() / \"badapple-small.mp4\"\n",
    "\n",
    "cap = cv2.VideoCapture(str(badapple_path))\n",
    "\n",
    "width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))\n",
    "height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))\n",
    "fps = int(cap.get(cv2.CAP_PROP_FPS))\n",
    "n_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))\n",
    "print(f\"width: {width}, height: {height}, fps: {fps}, n_frames: {n_frames}\")\n",
    "\n",
    "square_side = 30\n",
    "black_threshold = 50"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-02-24T18:40:45.231723500Z",
     "start_time": "2024-02-24T18:40:44.737443600Z"
    }
   },
   "id": "5dc3eefb4ee47cc2"
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "outputs": [],
   "source": [
    "circle_frames_dir = Path(\"circle_frames\")\n",
    "circle_frames_dir.mkdir(exist_ok=True)\n",
    "real_frames_dir = Path(\"real_frames\")\n",
    "real_frames_dir.mkdir(exist_ok=True)\n",
    "diff_frames_dir = Path(\"diff_frames\")\n",
    "diff_frames_dir.mkdir(exist_ok=True)\n",
    "diff_frames_tsp_dir = Path(\"diff_frames_tsp\")\n",
    "diff_frames_tsp_dir.mkdir(exist_ok=True)"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-02-24T18:40:46.508949Z",
     "start_time": "2024-02-24T18:40:46.496573200Z"
    }
   },
   "id": "ab27e2f2883c772f"
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Store real_frames from the original video\n"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "9d7b64bc69a76ece"
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "frame 6572/6572\r"
     ]
    }
   ],
   "source": [
    "ret, frame = cap.read()\n",
    "index = 0\n",
    "while ret:\n",
    "    cv2.imwrite(f\"real_frames/{index}.jpg\", frame)\n",
    "    ret, frame = cap.read()\n",
    "    index += 1\n",
    "    print(f\"frame {index}/{n_frames}\", end=\"\\r\")"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-02-03T04:07:34.785923700Z",
     "start_time": "2024-02-03T04:06:47.048289300Z"
    }
   },
   "id": "1ad2fd8035277a66"
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Generate circle, downsampled frames with coordinates of apples labeled"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "564f8cf3743ad5aa"
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "frame 6571/6572\r"
     ]
    }
   ],
   "source": [
    "# reset cap\n",
    "cap = cv2.VideoCapture(str(badapple_path))\n",
    "index = 0\n",
    "while cap.isOpened():\n",
    "    ret, frame = cap.read()\n",
    "    if not ret:\n",
    "        break\n",
    "    # convert to grayscale\n",
    "    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)\n",
    "    if index % 100 == 0 or index == n_frames - 1:\n",
    "        print(f\"frame {index}/{n_frames}\", end=\"\\r\")\n",
    "    out_frame = np.ones((height, width), dtype=np.uint8) * 255\n",
    "\n",
    "    current_frame_pixels = set()\n",
    "    for i in range(0, width - 1, square_side):\n",
    "        for j in range(0, height - 1, square_side):\n",
    "            section = frame[j: j + square_side, i: i + square_side]\n",
    "            # if its dark, add dark circle with white coordinate text\n",
    "            if (np.mean(section)) < black_threshold:\n",
    "                cv2.circle(out_frame, (i+square_side//2, j+square_side//2), square_side//2, 0, -1)\n",
    "                text_color = 255\n",
    "            else:\n",
    "                text_color = 0\n",
    "            # otherwise, black coordinate text\n",
    "            # put first coordinate higher, second coordinate lower\n",
    "            \n",
    "            cv2.putText(out_frame, f\"{j//square_side+1}\", (i+square_side//3-1, j + 22), cv2.FONT_HERSHEY_SIMPLEX, 0.3, text_color, 1, cv2.LINE_AA)\n",
    "            cv2.putText(out_frame, f\"{i//square_side+1}\", (i+square_side//3-1, j + 12), cv2.FONT_HERSHEY_SIMPLEX, 0.3, text_color, 1, cv2.LINE_AA)\n",
    "    cv2.imwrite(f\"circle_frames/{index}.jpg\", out_frame)\n",
    "    index += 1"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-02-03T04:09:29.264262100Z",
     "start_time": "2024-02-03T04:07:37.042919500Z"
    }
   },
   "id": "a3a646311b9b3e6d"
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Generate diff frames with additions in green and removals in red"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "5902751fd37754b7"
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "frame 6571/6572\r"
     ]
    }
   ],
   "source": [
    "# reset cap\n",
    "cap = cv2.VideoCapture(str(badapple_path))\n",
    "index = 0\n",
    "last_frame_pixels = set()\n",
    "\n",
    "while cap.isOpened():\n",
    "    ret, frame = cap.read()\n",
    "\n",
    "    if not ret:\n",
    "        break\n",
    "    # convert to grayscale\n",
    "    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)\n",
    "    if index % 100 == 0 or index == n_frames - 1:\n",
    "        print(f\"frame {index}/{n_frames}\", end=\"\\r\")\n",
    "        \n",
    "    out_frame = np.ones((height, width, 3), dtype=np.uint8) * 255\n",
    "\n",
    "    current_frame_pixels = set()\n",
    "    for i in range(0, width - 1, square_side):\n",
    "        for j in range(0, height - 1, square_side):\n",
    "            section = frame[j: j + square_side, i: i + square_side]\n",
    "             # three cases we care about: a pixel is persisted, a pixel is added, a pixel is removed\n",
    "            circle_colour = None\n",
    "            if (np.mean(section)) < 50:\n",
    "                current_frame_pixels.add((i, j))\n",
    "                # black pixel\n",
    "                if (i, j) in last_frame_pixels:\n",
    "                    # persisted from previous frame. Black circle, white text\n",
    "                    circle_colour = (0, 0, 0)\n",
    "                    text_color = (255, 255, 255)\n",
    "                else:\n",
    "                    # new black dot. Green circle, black text\n",
    "                    circle_colour = (0, 255, 0)\n",
    "                    text_color = (0, 0 ,0)\n",
    "            else:\n",
    "                # white pixel\n",
    "                if (i, j) in last_frame_pixels:\n",
    "                    # removed black dot. Red circle, black text\n",
    "                    circle_colour = (0, 0, 255)\n",
    "                    text_color = (0, 0, 0)\n",
    "                else:\n",
    "                    # persisted nothing. Just use black text\n",
    "                    text_color = (0, 0 ,0)\n",
    "                    \n",
    "            if circle_colour is not None:\n",
    "                cv2.circle(out_frame, (i+square_side//2, j+square_side//2), square_side//2, circle_colour, -1)\n",
    "            cv2.putText(out_frame, f\"{j//square_side+1}\", (i+square_side//3-1, j + 22), cv2.FONT_HERSHEY_SIMPLEX, 0.3, text_color, 1, cv2.LINE_AA)\n",
    "            cv2.putText(out_frame, f\"{i//square_side+1}\", (i+square_side//3-1, j + 12), cv2.FONT_HERSHEY_SIMPLEX, 0.3, text_color, 1, cv2.LINE_AA)\n",
    "    cv2.imwrite(f\"diff_frames/{index}.jpg\", out_frame)\n",
    "\n",
    "    last_frame_pixels = current_frame_pixels\n",
    "    index += 1"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-02-24T18:43:34.279733400Z",
     "start_time": "2024-02-24T18:40:54.866672Z"
    }
   },
   "id": "996035405a9ef9ef"
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Additions and removals in different colours but also arrows using TSP algorithm"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "411aed6a52b279e0"
  },
  {
   "cell_type": "markdown",
   "source": [
    "OK ignore this section, this optimization did not work and isn't needed, just move the apples using your own mental heuristics"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "21b8806df00cbfe2"
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "outputs": [],
   "source": [
    "# travelling salesman solving function\n",
    "from python_tsp.heuristics import solve_tsp_simulated_annealing\n",
    "from python_tsp.heuristics import solve_tsp_local_search\n",
    "from python_tsp.exact import solve_tsp_dynamic_programming\n",
    "from python_tsp.distances import euclidean_distance_matrix\n",
    "import numpy as np\n",
    "\n",
    "def closest_to_origin_index(points):\n",
    "    # given a list of points, return the point closest to (0, 0)\n",
    "    points_arr = np.array(points)\n",
    "    closest_to_origin = np.argmin(np.linalg.norm(points_arr, axis=1))\n",
    "    return closest_to_origin\n",
    "\n",
    "def tsp_solver(points):\n",
    "    # given a list of points, return a list of location tuples representing the path to take\n",
    "    # points is a list of (x, y) tuples, this is done on cartesian grid\n",
    "    # start closest to top right, because laptop will be in bottom right so top left is the least convenient to end on\n",
    "    points_arr = np.array(points)\n",
    "    # sort points_arr by distance from the origin actually\n",
    "    points_arr = points_arr[np.argsort(np.linalg.norm(points_arr, axis=1))]\n",
    "    # return points_arr\n",
    "    # find the path from that point to the rest of the points\n",
    "    # make nxn distance matrix\n",
    "    num_points = len(points)\n",
    "    distance_matrix = euclidean_distance_matrix(points_arr)\n",
    "    # make distance array where the distances from closest_to_origin to the rest of the points are 0\n",
    "    distance_matrix[0, :] = 0\n",
    "    # permutation, distance = solve_tsp_local_search(distance_matrix)\n",
    "    permutation, distance = solve_tsp_dynamic_programming(distance_matrix)\n",
    "    route = []\n",
    "    for i in range(num_points):\n",
    "        route.append(points[permutation[i]])\n",
    "    return route"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-12-29T08:09:31.384394Z",
     "start_time": "2023-12-29T08:09:31.365071Z"
    }
   },
   "id": "45a7d42ef3851c5a"
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "frame 0/6572, change_length: 768\n",
      "frame 1/6572, change_length: 0\n",
      "frame 2/6572, change_length: 0\n",
      "frame 3/6572, change_length: 0\n",
      "frame 4/6572, change_length: 0\n",
      "frame 5/6572, change_length: 0\n",
      "frame 6/6572, change_length: 0\n",
      "frame 7/6572, change_length: 0\n",
      "frame 8/6572, change_length: 0\n",
      "frame 9/6572, change_length: 0\n",
      "frame 10/6572, change_length: 0\n",
      "frame 11/6572, change_length: 0\n",
      "frame 12/6572, change_length: 0\n",
      "frame 13/6572, change_length: 0\n",
      "frame 14/6572, change_length: 0\n",
      "frame 15/6572, change_length: 0\n",
      "frame 16/6572, change_length: 0\n",
      "frame 17/6572, change_length: 0\n",
      "frame 18/6572, change_length: 0\n",
      "frame 19/6572, change_length: 0\n",
      "frame 20/6572, change_length: 0\n",
      "frame 21/6572, change_length: 0\n",
      "frame 22/6572, change_length: 0\n",
      "frame 23/6572, change_length: 0\n",
      "frame 24/6572, change_length: 0\n",
      "frame 25/6572, change_length: 0\n",
      "frame 26/6572, change_length: 0\n",
      "frame 27/6572, change_length: 0\n",
      "frame 28/6572, change_length: 0\n",
      "frame 29/6572, change_length: 0\n",
      "frame 30/6572, change_length: 0\n",
      "frame 31/6572, change_length: 0\n",
      "frame 32/6572, change_length: 0\n",
      "frame 33/6572, change_length: 0\n",
      "frame 34/6572, change_length: 0\n",
      "frame 35/6572, change_length: 0\n",
      "frame 36/6572, change_length: 0\n",
      "frame 37/6572, change_length: 0\n",
      "frame 38/6572, change_length: 0\n",
      "frame 39/6572, change_length: 0\n",
      "frame 40/6572, change_length: 0\n",
      "frame 41/6572, change_length: 2\n",
      "frame 42/6572, change_length: 22\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001B[0;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[0;31mKeyboardInterrupt\u001B[0m                         Traceback (most recent call last)",
      "Cell \u001B[0;32mIn[53], line 64\u001B[0m\n\u001B[1;32m     62\u001B[0m     cv2\u001B[38;5;241m.\u001B[39mputText(out_frame, \u001B[38;5;124mf\u001B[39m\u001B[38;5;124m\"\u001B[39m\u001B[38;5;124m2\u001B[39m\u001B[38;5;124m\"\u001B[39m, (change_arr[\u001B[38;5;241m1\u001B[39m\u001B[38;5;241m-\u001B[39mstart][\u001B[38;5;241m0\u001B[39m]\u001B[38;5;241m+\u001B[39msquare_side\u001B[38;5;241m/\u001B[39m\u001B[38;5;241m/\u001B[39m\u001B[38;5;241m3\u001B[39m\u001B[38;5;241m-\u001B[39m\u001B[38;5;241m1\u001B[39m, change_arr[\u001B[38;5;241m1\u001B[39m\u001B[38;5;241m-\u001B[39mstart][\u001B[38;5;241m1\u001B[39m] \u001B[38;5;241m+\u001B[39m \u001B[38;5;241m18\u001B[39m), cv2\u001B[38;5;241m.\u001B[39mFONT_HERSHEY_SIMPLEX, \u001B[38;5;241m0.3\u001B[39m, (\u001B[38;5;241m0\u001B[39m, \u001B[38;5;241m0\u001B[39m, \u001B[38;5;241m0\u001B[39m), \u001B[38;5;241m1\u001B[39m, cv2\u001B[38;5;241m.\u001B[39mLINE_AA)\n\u001B[1;32m     63\u001B[0m \u001B[38;5;28;01melif\u001B[39;00m index \u001B[38;5;241m!=\u001B[39m \u001B[38;5;241m0\u001B[39m \u001B[38;5;129;01mand\u001B[39;00m \u001B[38;5;28mlen\u001B[39m(pixels_that_change) \u001B[38;5;241m!=\u001B[39m \u001B[38;5;241m0\u001B[39m:\n\u001B[0;32m---> 64\u001B[0m     route \u001B[38;5;241m=\u001B[39m \u001B[43mtsp_solver\u001B[49m\u001B[43m(\u001B[49m\u001B[43mchange_list\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m     65\u001B[0m     \u001B[38;5;28;01mfor\u001B[39;00m i, point \u001B[38;5;129;01min\u001B[39;00m \u001B[38;5;28menumerate\u001B[39m(route):\n\u001B[1;32m     66\u001B[0m         cv2\u001B[38;5;241m.\u001B[39mputText(out_frame, \u001B[38;5;124mf\u001B[39m\u001B[38;5;124m\"\u001B[39m\u001B[38;5;132;01m{\u001B[39;00mi\u001B[38;5;241m+\u001B[39m\u001B[38;5;241m1\u001B[39m\u001B[38;5;132;01m}\u001B[39;00m\u001B[38;5;124m\"\u001B[39m, (point[\u001B[38;5;241m0\u001B[39m]\u001B[38;5;241m+\u001B[39msquare_side\u001B[38;5;241m/\u001B[39m\u001B[38;5;241m/\u001B[39m\u001B[38;5;241m3\u001B[39m\u001B[38;5;241m-\u001B[39m\u001B[38;5;241m1\u001B[39m, point[\u001B[38;5;241m1\u001B[39m] \u001B[38;5;241m+\u001B[39m \u001B[38;5;241m18\u001B[39m), cv2\u001B[38;5;241m.\u001B[39mFONT_HERSHEY_SIMPLEX, \u001B[38;5;241m0.3\u001B[39m, (\u001B[38;5;241m0\u001B[39m, \u001B[38;5;241m0\u001B[39m, \u001B[38;5;241m0\u001B[39m), \u001B[38;5;241m1\u001B[39m, cv2\u001B[38;5;241m.\u001B[39mLINE_AA)\n",
      "Cell \u001B[0;32mIn[52], line 29\u001B[0m, in \u001B[0;36mtsp_solver\u001B[0;34m(points)\u001B[0m\n\u001B[1;32m     27\u001B[0m distance_matrix[\u001B[38;5;241m0\u001B[39m, :] \u001B[38;5;241m=\u001B[39m \u001B[38;5;241m0\u001B[39m\n\u001B[1;32m     28\u001B[0m \u001B[38;5;66;03m# permutation, distance = solve_tsp_local_search(distance_matrix)\u001B[39;00m\n\u001B[0;32m---> 29\u001B[0m permutation, distance \u001B[38;5;241m=\u001B[39m \u001B[43msolve_tsp_dynamic_programming\u001B[49m\u001B[43m(\u001B[49m\u001B[43mdistance_matrix\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m     30\u001B[0m route \u001B[38;5;241m=\u001B[39m []\n\u001B[1;32m     31\u001B[0m \u001B[38;5;28;01mfor\u001B[39;00m i \u001B[38;5;129;01min\u001B[39;00m \u001B[38;5;28mrange\u001B[39m(num_points):\n",
      "File \u001B[0;32m~/src/badapple/venv/lib/python3.9/site-packages/python_tsp/exact/dynamic_programming.py:113\u001B[0m, in \u001B[0;36msolve_tsp_dynamic_programming\u001B[0;34m(distance_matrix, maxsize)\u001B[0m\n\u001B[1;32m    109\u001B[0m     memo[(ni, N)] \u001B[38;5;241m=\u001B[39m nmin\n\u001B[1;32m    111\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m min_cost\n\u001B[0;32m--> 113\u001B[0m best_distance \u001B[38;5;241m=\u001B[39m \u001B[43mdist\u001B[49m\u001B[43m(\u001B[49m\u001B[38;5;241;43m0\u001B[39;49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mN\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m    115\u001B[0m \u001B[38;5;66;03m# Step 2: get path with the minimum distance\u001B[39;00m\n\u001B[1;32m    116\u001B[0m ni \u001B[38;5;241m=\u001B[39m \u001B[38;5;241m0\u001B[39m  \u001B[38;5;66;03m# start at the origin\u001B[39;00m\n",
      "File \u001B[0;32m~/src/badapple/venv/lib/python3.9/site-packages/python_tsp/exact/dynamic_programming.py:104\u001B[0m, in \u001B[0;36msolve_tsp_dynamic_programming.<locals>.dist\u001B[0;34m(ni, N)\u001B[0m\n\u001B[1;32m    101\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m distance_matrix[ni, \u001B[38;5;241m0\u001B[39m]\n\u001B[1;32m    103\u001B[0m \u001B[38;5;66;03m# Store the costs in the form (nj, dist(nj, N))\u001B[39;00m\n\u001B[0;32m--> 104\u001B[0m costs \u001B[38;5;241m=\u001B[39m [\n\u001B[1;32m    105\u001B[0m     (nj, distance_matrix[ni, nj] \u001B[38;5;241m+\u001B[39m dist(nj, N\u001B[38;5;241m.\u001B[39mdifference({nj})))\n\u001B[1;32m    106\u001B[0m     \u001B[38;5;28;01mfor\u001B[39;00m nj \u001B[38;5;129;01min\u001B[39;00m N\n\u001B[1;32m    107\u001B[0m ]\n\u001B[1;32m    108\u001B[0m nmin, min_cost \u001B[38;5;241m=\u001B[39m \u001B[38;5;28mmin\u001B[39m(costs, key\u001B[38;5;241m=\u001B[39m\u001B[38;5;28;01mlambda\u001B[39;00m x: x[\u001B[38;5;241m1\u001B[39m])\n\u001B[1;32m    109\u001B[0m memo[(ni, N)] \u001B[38;5;241m=\u001B[39m nmin\n",
      "File \u001B[0;32m~/src/badapple/venv/lib/python3.9/site-packages/python_tsp/exact/dynamic_programming.py:105\u001B[0m, in \u001B[0;36m<listcomp>\u001B[0;34m(.0)\u001B[0m\n\u001B[1;32m    101\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m distance_matrix[ni, \u001B[38;5;241m0\u001B[39m]\n\u001B[1;32m    103\u001B[0m \u001B[38;5;66;03m# Store the costs in the form (nj, dist(nj, N))\u001B[39;00m\n\u001B[1;32m    104\u001B[0m costs \u001B[38;5;241m=\u001B[39m [\n\u001B[0;32m--> 105\u001B[0m     (nj, distance_matrix[ni, nj] \u001B[38;5;241m+\u001B[39m \u001B[43mdist\u001B[49m\u001B[43m(\u001B[49m\u001B[43mnj\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mN\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mdifference\u001B[49m\u001B[43m(\u001B[49m\u001B[43m{\u001B[49m\u001B[43mnj\u001B[49m\u001B[43m}\u001B[49m\u001B[43m)\u001B[49m\u001B[43m)\u001B[49m)\n\u001B[1;32m    106\u001B[0m     \u001B[38;5;28;01mfor\u001B[39;00m nj \u001B[38;5;129;01min\u001B[39;00m N\n\u001B[1;32m    107\u001B[0m ]\n\u001B[1;32m    108\u001B[0m nmin, min_cost \u001B[38;5;241m=\u001B[39m \u001B[38;5;28mmin\u001B[39m(costs, key\u001B[38;5;241m=\u001B[39m\u001B[38;5;28;01mlambda\u001B[39;00m x: x[\u001B[38;5;241m1\u001B[39m])\n\u001B[1;32m    109\u001B[0m memo[(ni, N)] \u001B[38;5;241m=\u001B[39m nmin\n",
      "File \u001B[0;32m~/src/badapple/venv/lib/python3.9/site-packages/python_tsp/exact/dynamic_programming.py:104\u001B[0m, in \u001B[0;36msolve_tsp_dynamic_programming.<locals>.dist\u001B[0;34m(ni, N)\u001B[0m\n\u001B[1;32m    101\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m distance_matrix[ni, \u001B[38;5;241m0\u001B[39m]\n\u001B[1;32m    103\u001B[0m \u001B[38;5;66;03m# Store the costs in the form (nj, dist(nj, N))\u001B[39;00m\n\u001B[0;32m--> 104\u001B[0m costs \u001B[38;5;241m=\u001B[39m [\n\u001B[1;32m    105\u001B[0m     (nj, distance_matrix[ni, nj] \u001B[38;5;241m+\u001B[39m dist(nj, N\u001B[38;5;241m.\u001B[39mdifference({nj})))\n\u001B[1;32m    106\u001B[0m     \u001B[38;5;28;01mfor\u001B[39;00m nj \u001B[38;5;129;01min\u001B[39;00m N\n\u001B[1;32m    107\u001B[0m ]\n\u001B[1;32m    108\u001B[0m nmin, min_cost \u001B[38;5;241m=\u001B[39m \u001B[38;5;28mmin\u001B[39m(costs, key\u001B[38;5;241m=\u001B[39m\u001B[38;5;28;01mlambda\u001B[39;00m x: x[\u001B[38;5;241m1\u001B[39m])\n\u001B[1;32m    109\u001B[0m memo[(ni, N)] \u001B[38;5;241m=\u001B[39m nmin\n",
      "File \u001B[0;32m~/src/badapple/venv/lib/python3.9/site-packages/python_tsp/exact/dynamic_programming.py:105\u001B[0m, in \u001B[0;36m<listcomp>\u001B[0;34m(.0)\u001B[0m\n\u001B[1;32m    101\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m distance_matrix[ni, \u001B[38;5;241m0\u001B[39m]\n\u001B[1;32m    103\u001B[0m \u001B[38;5;66;03m# Store the costs in the form (nj, dist(nj, N))\u001B[39;00m\n\u001B[1;32m    104\u001B[0m costs \u001B[38;5;241m=\u001B[39m [\n\u001B[0;32m--> 105\u001B[0m     (nj, distance_matrix[ni, nj] \u001B[38;5;241m+\u001B[39m \u001B[43mdist\u001B[49m\u001B[43m(\u001B[49m\u001B[43mnj\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mN\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mdifference\u001B[49m\u001B[43m(\u001B[49m\u001B[43m{\u001B[49m\u001B[43mnj\u001B[49m\u001B[43m}\u001B[49m\u001B[43m)\u001B[49m\u001B[43m)\u001B[49m)\n\u001B[1;32m    106\u001B[0m     \u001B[38;5;28;01mfor\u001B[39;00m nj \u001B[38;5;129;01min\u001B[39;00m N\n\u001B[1;32m    107\u001B[0m ]\n\u001B[1;32m    108\u001B[0m nmin, min_cost \u001B[38;5;241m=\u001B[39m \u001B[38;5;28mmin\u001B[39m(costs, key\u001B[38;5;241m=\u001B[39m\u001B[38;5;28;01mlambda\u001B[39;00m x: x[\u001B[38;5;241m1\u001B[39m])\n\u001B[1;32m    109\u001B[0m memo[(ni, N)] \u001B[38;5;241m=\u001B[39m nmin\n",
      "    \u001B[0;31m[... skipping similar frames: <listcomp> at line 105 (10 times), solve_tsp_dynamic_programming.<locals>.dist at line 104 (10 times)]\u001B[0m\n",
      "File \u001B[0;32m~/src/badapple/venv/lib/python3.9/site-packages/python_tsp/exact/dynamic_programming.py:104\u001B[0m, in \u001B[0;36msolve_tsp_dynamic_programming.<locals>.dist\u001B[0;34m(ni, N)\u001B[0m\n\u001B[1;32m    101\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m distance_matrix[ni, \u001B[38;5;241m0\u001B[39m]\n\u001B[1;32m    103\u001B[0m \u001B[38;5;66;03m# Store the costs in the form (nj, dist(nj, N))\u001B[39;00m\n\u001B[0;32m--> 104\u001B[0m costs \u001B[38;5;241m=\u001B[39m [\n\u001B[1;32m    105\u001B[0m     (nj, distance_matrix[ni, nj] \u001B[38;5;241m+\u001B[39m dist(nj, N\u001B[38;5;241m.\u001B[39mdifference({nj})))\n\u001B[1;32m    106\u001B[0m     \u001B[38;5;28;01mfor\u001B[39;00m nj \u001B[38;5;129;01min\u001B[39;00m N\n\u001B[1;32m    107\u001B[0m ]\n\u001B[1;32m    108\u001B[0m nmin, min_cost \u001B[38;5;241m=\u001B[39m \u001B[38;5;28mmin\u001B[39m(costs, key\u001B[38;5;241m=\u001B[39m\u001B[38;5;28;01mlambda\u001B[39;00m x: x[\u001B[38;5;241m1\u001B[39m])\n\u001B[1;32m    109\u001B[0m memo[(ni, N)] \u001B[38;5;241m=\u001B[39m nmin\n",
      "File \u001B[0;32m~/src/badapple/venv/lib/python3.9/site-packages/python_tsp/exact/dynamic_programming.py:105\u001B[0m, in \u001B[0;36m<listcomp>\u001B[0;34m(.0)\u001B[0m\n\u001B[1;32m    101\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m distance_matrix[ni, \u001B[38;5;241m0\u001B[39m]\n\u001B[1;32m    103\u001B[0m \u001B[38;5;66;03m# Store the costs in the form (nj, dist(nj, N))\u001B[39;00m\n\u001B[1;32m    104\u001B[0m costs \u001B[38;5;241m=\u001B[39m [\n\u001B[0;32m--> 105\u001B[0m     (nj, distance_matrix[ni, nj] \u001B[38;5;241m+\u001B[39m \u001B[43mdist\u001B[49m\u001B[43m(\u001B[49m\u001B[43mnj\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mN\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mdifference\u001B[49m\u001B[43m(\u001B[49m\u001B[43m{\u001B[49m\u001B[43mnj\u001B[49m\u001B[43m}\u001B[49m\u001B[43m)\u001B[49m\u001B[43m)\u001B[49m)\n\u001B[1;32m    106\u001B[0m     \u001B[38;5;28;01mfor\u001B[39;00m nj \u001B[38;5;129;01min\u001B[39;00m N\n\u001B[1;32m    107\u001B[0m ]\n\u001B[1;32m    108\u001B[0m nmin, min_cost \u001B[38;5;241m=\u001B[39m \u001B[38;5;28mmin\u001B[39m(costs, key\u001B[38;5;241m=\u001B[39m\u001B[38;5;28;01mlambda\u001B[39;00m x: x[\u001B[38;5;241m1\u001B[39m])\n\u001B[1;32m    109\u001B[0m memo[(ni, N)] \u001B[38;5;241m=\u001B[39m nmin\n",
      "\u001B[0;31mKeyboardInterrupt\u001B[0m: "
     ]
    }
   ],
   "source": [
    "# reset cap\n",
    "cap = cv2.VideoCapture(str(badapple_path))\n",
    "index = 0\n",
    "last_frame_pixels = set()\n",
    "\n",
    "while cap.isOpened():\n",
    "    ret, frame = cap.read()\n",
    "\n",
    "    if not ret:\n",
    "        break\n",
    "    # convert to grayscale\n",
    "    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)\n",
    "    if index % 100 == 0 or index == n_frames - 1:\n",
    "        print(f\"frame {index}/{n_frames}\", end=\"\\r\")\n",
    "        \n",
    "    out_frame = np.ones((height, width, 3), dtype=np.uint8) * 255\n",
    "\n",
    "    current_frame_pixels = set()\n",
    "    pixels_that_change = set()\n",
    "    for i in range(0, width - 1, square_side):\n",
    "        for j in range(0, height - 1, square_side):\n",
    "            section = frame[j: j + square_side, i: i + square_side]\n",
    "             # three cases we care about: a pixel is persisted, a pixel is added, a pixel is removed\n",
    "            text_color = None\n",
    "            circle_colour = None\n",
    "            if (np.mean(section)) < 50:\n",
    "                current_frame_pixels.add((i, j))\n",
    "                # black pixel\n",
    "                if (i, j) in last_frame_pixels:\n",
    "                    # persisted from previous frame. Black circle, white text\n",
    "                    circle_colour = (0, 0, 0)\n",
    "                    text_color = (255, 255, 255)\n",
    "                else:\n",
    "                    # new black dot. Green circle, no text because tsp markers will be used\n",
    "                    circle_colour = (0, 255, 0)\n",
    "                    pixels_that_change.add((i, j))\n",
    "            else:\n",
    "                # white pixel\n",
    "                if (i, j) in last_frame_pixels:\n",
    "                    # removed black dot. Red circle, no text because tsp markers will be used\n",
    "                    circle_colour = (0, 0, 255)\n",
    "                    pixels_that_change.add((i, j))\n",
    "                else:\n",
    "                    # persisted nothing. Just use black text\n",
    "                    text_color = (0, 0 ,0)\n",
    "                    \n",
    "            if circle_colour is not None:\n",
    "                cv2.circle(out_frame, (i+square_side//2, j+square_side//2), square_side//2, circle_colour, -1)\n",
    "            if text_color is not None:\n",
    "                cv2.putText(out_frame, f\"{j//square_side+1}\", (i+square_side//3-1, j + 22), cv2.FONT_HERSHEY_SIMPLEX, 0.3, text_color, 1, cv2.LINE_AA)\n",
    "                cv2.putText(out_frame, f\"{i//square_side+1}\", (i+square_side//3-1, j + 12), cv2.FONT_HERSHEY_SIMPLEX, 0.3, text_color, 1, cv2.LINE_AA)\n",
    "    change_list = list(pixels_that_change)\n",
    "    change_length = len(change_list)\n",
    "    print(f\"frame {index}/{n_frames}, change_length: {change_length}\")\n",
    "    if change_length == 1:\n",
    "        # one change, no arrows\n",
    "        pass\n",
    "    elif change_length == 2: # only two changes, don't run TSP, label the closest one 1 and the other 2\n",
    "        start = closest_to_origin_index(change_list)\n",
    "        change_arr = np.array(change_list) # only used so I can add square_side//2 to the coordinates without having to index the tuple\n",
    "        cv2.putText(out_frame, f\"1\", (change_arr[start][0]+square_side//3-1, change_arr[start][1] + 18), cv2.FONT_HERSHEY_SIMPLEX, 0.3, (0, 0, 0), 1, cv2.LINE_AA)\n",
    "        cv2.putText(out_frame, f\"2\", (change_arr[1-start][0]+square_side//3-1, change_arr[1-start][1] + 18), cv2.FONT_HERSHEY_SIMPLEX, 0.3, (0, 0, 0), 1, cv2.LINE_AA)\n",
    "    elif index != 0 and len(pixels_that_change) != 0:\n",
    "        route = tsp_solver(change_list)\n",
    "        for i, point in enumerate(route):\n",
    "            cv2.putText(out_frame, f\"{i+1}\", (point[0]+square_side//3-1, point[1] + 18), cv2.FONT_HERSHEY_SIMPLEX, 0.3, (0, 0, 0), 1, cv2.LINE_AA)\n",
    "            \n",
    "    cv2.imwrite(f\"diff_frames_tsp/{index}.jpg\", out_frame)\n",
    "\n",
    "    last_frame_pixels = current_frame_pixels\n",
    "    index += 1"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-12-29T08:16:24.402334Z",
     "start_time": "2023-12-29T08:09:32.549267Z"
    }
   },
   "id": "1c2e314a004f8bfb"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [],
   "metadata": {
    "collapsed": false
   },
   "id": "45db082f4c03483f"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}