{ "cells": [ { "cell_type": "markdown", "id": "e4105dae", "metadata": {}, "source": [ "# Spin Test 3" ] }, { "cell_type": "code", "execution_count": 13, "id": "e6c3f608", "metadata": {}, "outputs": [], "source": [ "import cv2\n", "from pathlib import Path\n", "import numpy as np\n", "import pandas as pd" ] }, { "cell_type": "markdown", "id": "d7baad73", "metadata": {}, "source": [ "Import video" ] }, { "cell_type": "code", "execution_count": 14, "id": "7521fd8e", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Opened: True, FPS: 59.94005994005994, Total Frames: 227.0\n" ] } ], "source": [ "video_number = \"3\"\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": "796d8a4d", "metadata": {}, "source": [ "Import data" ] }, { "cell_type": "code", "execution_count": 15, "id": "c25ccb49", "metadata": {}, "outputs": [], "source": [ "# Define the path to the CSV file\n", "# input_data_path = project_root / \"notebook\" / \"ball_detection\" / \"intermediate_data\" / f\"Circle_positions_cleaned_{video_number}.csv\"\n", "input_data_path = (\n", " project_root\n", " / \"data\"\n", " / \"auxiliary_data\"\n", " / \"circle_positions\"\n", " / f\"Adjusted_positions_new_{video_number}.csv\"\n", ")\n", "\n", "# Load the CSV file into a DataFrame\n", "df = pd.read_csv(input_data_path)" ] }, { "cell_type": "code", "execution_count": null, "id": "00a9afdc", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Frame 37: Found 7 points to track\n", "Frame 37: Found 7 points to track\n", "Frame 37: Found 7 points to track\n", "Frame 38: Found 4 points to track\n", "Frame 38: Found 4 points to track\n", "Frame 38: Found 4 points to track\n", "Frame 39: Found 9 points to track\n", "Frame 39: Found 9 points to track\n", "Frame 39: Found 9 points to track\n", "Frame 40: Found 9 points to track\n", "Frame 40: Found 9 points to track\n", "Frame 40: Found 9 points to track\n", "Frame 41: Found 4 points to track\n", "Frame 41: Found 4 points to track\n", "Frame 41: Found 4 points to track\n", "Frame 42: Found 5 points to track\n", "Frame 42: Found 5 points to track\n", "Frame 42: Found 5 points to track\n", "Frame 43: Found 7 points to track\n", "Frame 43: Found 7 points to track\n", "Frame 43: Found 7 points to track\n", "Frame 44: Found 5 points to track\n", "Frame 44: Found 5 points to track\n", "Frame 44: Found 5 points to track\n", "Frame 45: Found 4 points to track\n", "Frame 45: Found 4 points to track\n", "Frame 45: Found 4 points to track\n", "Frame 46: Found 7 points to track\n", "Frame 46: Found 7 points to track\n", "Frame 46: Found 7 points to track\n", "Frame 47: Found 2 points to track\n", "Frame 47: Found 2 points to track\n", "Frame 47: Found 2 points to track\n", "Frame 48: Found 4 points to track\n", "Frame 48: Found 4 points to track\n", "Frame 48: Found 4 points to track\n", "Frame 49: Found 2 points to track\n", "Frame 49: Found 2 points to track\n", "Frame 49: Found 2 points to track\n", "Frame 50: Found 5 points to track\n", "Frame 50: Found 5 points to track\n", "Frame 50: Found 5 points to track\n", "Frame 51: Found 1 points to track\n", "Frame 51: Found 1 points to track\n", "Frame 51: Found 1 points to track\n", "Frame 52: Found 7 points to track\n", "Frame 52: Found 7 points to track\n", "Frame 52: Found 7 points to track\n", "Frame 53: Found 5 points to track\n", "Frame 53: Found 5 points to track\n", "Frame 53: Found 5 points to track\n", "Frame 54: Found 5 points to track\n", "Frame 54: Found 5 points to track\n", "Frame 54: Found 5 points to track\n", "Frame 55: Found 1 points to track\n", "Frame 55: Found 1 points to track\n", "Frame 55: Found 1 points to track\n", "Frame 56: Found 1 points to track\n", "Frame 56: Found 1 points to track\n", "Frame 56: Found 1 points to track\n", "Frame 57: Found 4 points to track\n", "Frame 57: Found 4 points to track\n", "Frame 57: Found 4 points to track\n", "Frame 58: Found 4 points to track\n", "Frame 58: Found 4 points to track\n", "Frame 58: Found 4 points to track\n", "Frame 59: Found 2 points to track\n", "Frame 59: Found 2 points to track\n", "Frame 59: Found 2 points to track\n", "Frame 60: Found 3 points to track\n", "Frame 60: Found 3 points to track\n", "Frame 60: Found 3 points to track\n", "Frame 61: Found 5 points to track\n", "Frame 61: Found 5 points to track\n", "Frame 61: Found 5 points to track\n", "Frame 62: Found 8 points to track\n", "Frame 62: Found 8 points to track\n", "Frame 62: Found 8 points to track\n", "Frame 63: Found 2 points to track\n", "Frame 63: Found 2 points to track\n", "Frame 63: Found 2 points to track\n", "Frame 64: Found 4 points to track\n", "Frame 64: Found 4 points to track\n", "Frame 64: Found 4 points to track\n", "Frame 65: Found 4 points to track\n", "Frame 65: Found 4 points to track\n", "Frame 65: Found 4 points to track\n", "Frame 66: Found 4 points to track\n", "Frame 66: Found 4 points to track\n", "Frame 66: Found 4 points to track\n", "Frame 67: Found 4 points to track\n", "Frame 67: Found 4 points to track\n", "Frame 67: Found 4 points to track\n", "Frame 68: Found 8 points to track\n", "Frame 68: Found 8 points to track\n", "Frame 68: Found 8 points to track\n", "Frame 69: Found 6 points to track\n", "Frame 69: Found 6 points to track\n", "Frame 69: Found 6 points to track\n", "Frame 70: Found 5 points to track\n", "Frame 70: Found 5 points to track\n", "Frame 70: Found 5 points to track\n", "Frame 71: Found 3 points to track\n", "Frame 71: Found 3 points to track\n", "Frame 71: Found 3 points to track\n", "Frame 72: Found 4 points to track\n", "Frame 72: Found 4 points to track\n", "Frame 72: Found 4 points to track\n", "Frame 73: Found 3 points to track\n", "Frame 73: Found 3 points to track\n", "Frame 73: Found 3 points to track\n", "Frame 74: Found 3 points to track\n", "Frame 74: Found 3 points to track\n", "Frame 74: Found 3 points to track\n", "Frame 75: Found 5 points to track\n", "Frame 75: Found 5 points to track\n", "Frame 75: Found 5 points to track\n", "Frame 76: Found 6 points to track\n", "Frame 76: Found 6 points to track\n", "Frame 76: Found 6 points to track\n", "Frame 77: Found 4 points to track\n", "Frame 77: Found 4 points to track\n", "Frame 77: Found 4 points to track\n", "Frame 78: Found 5 points to track\n", "Frame 78: Found 5 points to track\n", "Frame 78: Found 5 points to track\n", "Frame 79: Found 4 points to track\n", "Frame 79: Found 4 points to track\n", "Frame 79: Found 4 points to track\n", "Frame 80: Found 4 points to track\n", "Frame 80: Found 4 points to track\n", "Frame 80: Found 4 points to track\n", "Frame 81: Found 3 points to track\n", "Frame 81: Found 3 points to track\n", "Frame 81: Found 3 points to track\n", "Frame 82: Found 4 points to track\n", "Frame 82: Found 4 points to track\n", "Frame 82: Found 4 points to track\n", "Frame 83: Found 5 points to track\n", "Frame 83: Found 5 points to track\n", "Frame 83: Found 5 points to track\n", "Frame 84: Found 7 points to track\n", "Frame 84: Found 7 points to track\n", "Frame 84: Found 7 points to track\n", "Frame 85: Found 7 points to track\n", "Frame 85: Found 7 points to track\n", "Frame 85: Found 7 points to track\n", "Frame 86: Found 5 points to track\n", "Frame 86: Found 5 points to track\n", "Frame 86: Found 5 points to track\n", "Frame 87: Found 2 points to track\n", "Frame 87: Found 2 points to track\n", "Frame 87: Found 2 points to track\n", "Frame 88: Found 2 points to track\n", "Frame 88: Found 2 points to track\n", "Frame 88: Found 2 points to track\n", "Frame 89: Found 3 points to track\n", "Frame 89: Found 3 points to track\n", "Frame 89: Found 3 points to track\n", "Frame 90: Found 4 points to track\n", "Frame 90: Found 4 points to track\n", "Frame 90: Found 4 points to track\n", "Frame 91: Found 3 points to track\n", "Frame 91: Found 3 points to track\n", "Frame 91: Found 3 points to track\n", "Frame 92: Found 2 points to track\n", "Frame 92: Found 2 points to track\n", "Frame 92: Found 2 points to track\n", "Frame 93: Found 2 points to track\n", "Frame 93: Found 2 points to track\n", "Frame 93: Found 2 points to track\n", "Frame 94: Found 5 points to track\n", "Frame 94: Found 5 points to track\n", "Frame 94: Found 5 points to track\n", "Frame 95: Found 3 points to track\n", "Frame 95: Found 3 points to track\n", "Frame 95: Found 3 points to track\n", "Frame 96: Found 2 points to track\n", "Frame 96: Found 2 points to track\n", "Frame 96: Found 2 points to track\n", "Frame 97: Found 3 points to track\n", "Frame 97: Found 3 points to track\n", "Frame 97: Found 3 points to track\n", "Frame 98: Found 2 points to track\n", "Frame 98: Found 2 points to track\n", "Frame 98: Found 2 points to track\n", "Frame 99: Found 3 points to track\n", "Frame 99: Found 3 points to track\n", "Frame 99: Found 3 points to track\n", "Frame 100: Found 4 points to track\n", "Frame 100: Found 4 points to track\n", "Frame 100: Found 4 points to track\n", "Frame 101: Found 3 points to track\n", "Frame 101: Found 3 points to track\n", "Frame 101: Found 3 points to track\n", "Frame 102: Found 4 points to track\n", "Frame 102: Found 4 points to track\n", "Frame 102: Found 4 points to track\n", "Frame 103: Found 3 points to track\n", "Frame 103: Found 3 points to track\n", "Frame 103: Found 3 points to track\n", "Frame 104: Found 1 points to track\n", "Frame 104: Found 1 points to track\n", "Frame 104: Found 1 points to track\n", "Frame 105: Found 4 points to track\n", "Frame 105: Found 4 points to track\n", "Frame 105: Found 4 points to track\n", "Frame 106: Found 1 points to track\n", "Frame 106: Found 1 points to track\n", "Frame 106: Found 1 points to track\n", "Frame 107: Found 2 points to track\n", "Frame 107: Found 2 points to track\n", "Frame 107: Found 2 points to track\n", "Frame 108: Found 2 points to track\n", "Frame 108: Found 2 points to track\n", "Frame 108: Found 2 points to track\n", "Frame 109: Found 4 points to track\n", "Frame 109: Found 4 points to track\n", "Frame 109: Found 4 points to track\n", "Frame 110: Found 4 points to track\n", "Frame 110: Found 4 points to track\n", "Frame 110: Found 4 points to track\n", "Frame 111: Found 3 points to track\n", "Frame 111: Found 3 points to track\n", "Frame 111: Found 3 points to track\n", "Frame 112: Found 2 points to track\n", "Frame 112: Found 2 points to track\n", "Frame 112: Found 2 points to track\n", "Frame 113: Found 2 points to track\n", "Frame 113: Found 2 points to track\n", "Frame 113: Found 2 points to track\n", "Frame 114: Found 1 points to track\n", "Frame 114: Found 1 points to track\n", "Frame 114: Found 1 points to track\n", "Frame 115: Found 2 points to track\n", "Frame 115: Found 2 points to track\n", "Frame 115: Found 2 points to track\n", "Frame 116: Found 1 points to track\n", "Frame 116: Found 1 points to track\n", "Frame 116: Found 1 points to track\n", "Frame 117: Found 2 points to track\n", "Frame 117: Found 2 points to track\n", "Frame 117: Found 2 points to track\n", "Frame 118: Found 2 points to track\n", "Frame 118: Found 2 points to track\n", "Frame 118: Found 2 points to track\n", "Frame 119: Found 3 points to track\n", "Frame 119: Found 3 points to track\n", "Frame 119: Found 3 points to track\n", "Frame 120: Found 4 points to track\n", "Frame 120: Found 4 points to track\n", "Frame 120: Found 4 points to track\n", "Frame 121: Found 5 points to track\n", "Frame 121: Found 5 points to track\n", "Frame 121: Found 5 points to track\n", "Frame 122: Found 1 points to track\n", "Frame 122: Found 1 points to track\n", "Frame 122: Found 1 points to track\n", "Frame 123: Found 2 points to track\n", "Frame 123: Found 2 points to track\n", "Frame 123: Found 2 points to track\n", "Frame 124: Found 1 points to track\n", "Frame 124: Found 1 points to track\n", "Frame 124: Found 1 points to track\n", "Frame 125: Found 2 points to track\n", "Frame 125: Found 2 points to track\n", "Frame 125: Found 2 points to track\n", "Frame 126: Found 2 points to track\n", "Frame 126: Found 2 points to track\n", "Frame 126: Found 2 points to track\n", "Frame 127: Found 3 points to track\n", "Frame 127: Found 3 points to track\n", "Frame 127: Found 3 points to track\n", "Frame 128: Found 3 points to track\n", "Frame 128: Found 3 points to track\n", "Frame 128: Found 3 points to track\n", "Frame 129: Found 15 points to track\n", "Frame 129: Found 15 points to track\n", "Frame 129: Found 15 points to track\n", "Frame 130: Found 4 points to track\n", "Frame 130: Found 4 points to track\n", "Frame 130: Found 4 points to track\n", "Frame 131: Found 2 points to track\n", "Frame 131: Found 2 points to track\n", "Frame 131: Found 2 points to track\n", "Frame 132: Found 1 points to track\n", "Frame 132: Found 1 points to track\n", "Frame 132: Found 1 points to track\n", "Frame 133: Found 5 points to track\n", "Frame 133: Found 5 points to track\n", "Frame 133: Found 5 points to track\n", "Frame 134: Found 4 points to track\n", "Frame 134: Found 4 points to track\n", "Frame 134: Found 4 points to track\n", "Frame 135: Found 7 points to track\n", "Frame 135: Found 7 points to track\n", "Frame 135: Found 7 points to track\n", "Frame 136: Found 1 points to track\n", "Frame 136: Found 1 points to track\n", "Frame 136: Found 1 points to track\n", "Frame 137: Found 8 points to track\n", "Frame 137: Found 8 points to track\n", "Frame 137: Found 8 points to track\n", "Frame 138: Found 2 points to track\n", "Frame 138: Found 2 points to track\n", "Frame 138: Found 2 points to track\n", "Frame 139: Found 3 points to track\n", "Frame 139: Found 3 points to track\n", "Frame 139: Found 3 points to track\n", "Frame 140: Found 11 points to track\n", "Frame 140: Found 11 points to track\n", "Frame 140: Found 11 points to track\n", "Frame 141: Found 1 points to track\n", "Frame 141: Found 1 points to track\n", "Frame 141: Found 1 points to track\n", "Frame 142: Found 2 points to track\n", "Frame 142: Found 2 points to track\n", "Frame 142: Found 2 points to track\n", "Frame 143: Found 6 points to track\n", "Frame 143: Found 6 points to track\n", "Frame 143: Found 6 points to track\n", "Frame 144: Found 5 points to track\n", "Frame 144: Found 5 points to track\n", "Frame 144: Found 5 points to track\n", "Frame 145: Found 9 points to track\n", "Frame 145: Found 9 points to track\n", "Frame 145: Found 9 points to track\n", "Frame 146: Found 2 points to track\n", "Frame 146: Found 2 points to track\n", "Frame 146: Found 2 points to track\n", "Frame 147: Found 1 points to track\n", "Frame 147: Found 1 points to track\n", "Frame 147: Found 1 points to track\n", "Frame 148: Found 4 points to track\n", "Frame 148: Found 4 points to track\n", "Frame 148: Found 4 points to track\n", "Frame 149: Found 2 points to track\n", "Frame 149: Found 2 points to track\n", "Frame 149: Found 2 points to track\n", "Frame 150: Found 5 points to track\n", "Frame 150: Found 5 points to track\n", "Frame 150: Found 5 points to track\n", "Frame 151: Found 3 points to track\n", "Frame 151: Found 3 points to track\n", "Frame 151: Found 3 points to track\n", "Frame 152: Found 2 points to track\n", "Frame 152: Found 2 points to track\n", "Frame 152: Found 2 points to track\n", "Frame 153: Found 3 points to track\n", "Frame 153: Found 3 points to track\n", "Frame 153: Found 3 points to track\n", "Frame 154: Found 4 points to track\n", "Frame 154: Found 4 points to track\n", "Frame 154: Found 4 points to track\n", "Frame 155: Found 7 points to track\n", "Frame 155: Found 7 points to track\n", "Frame 155: Found 7 points to track\n", "Frame 156: Found 6 points to track\n", "Frame 156: Found 6 points to track\n", "Frame 156: Found 6 points to track\n", "Frame 157: Found 6 points to track\n", "Frame 157: Found 6 points to track\n", "Frame 157: Found 6 points to track\n", "Frame 158: Found 7 points to track\n", "Frame 158: Found 7 points to track\n", "Frame 158: Found 7 points to track\n", "Frame 159: Found 9 points to track\n", "Frame 159: Found 9 points to track\n", "Frame 159: Found 9 points to track\n", "Frame 160: Found 5 points to track\n", "Frame 160: Found 5 points to track\n", "Frame 160: Found 5 points to track\n", "Frame 161: Found 9 points to track\n", "Frame 161: Found 9 points to track\n", "Frame 161: Found 9 points to track\n", "Frame 162: Found 12 points to track\n", "Frame 162: Found 12 points to track\n", "Frame 162: Found 12 points to track\n", "Frame 163: Found 9 points to track\n", "Frame 163: Found 9 points to track\n", "Frame 163: Found 9 points to track\n", "Frame 164: Found 7 points to track\n", "Frame 164: Found 7 points to track\n", "Frame 164: Found 7 points to track\n", "Frame 165: Found 3 points to track\n", "Frame 165: Found 3 points to track\n", "Frame 165: Found 3 points to track\n", "Frame 166: Found 4 points to track\n", "Frame 166: Found 4 points to track\n", "Frame 166: Found 4 points to track\n", "Frame 167: Found 7 points to track\n", "Frame 167: Found 7 points to track\n", "Frame 167: Found 7 points to track\n", "Frame 168: Found 6 points to track\n", "Frame 168: Found 6 points to track\n", "Frame 168: Found 6 points to track\n", "Frame 169: Found 5 points to track\n", "Frame 169: Found 5 points to track\n", "Frame 169: Found 5 points to track\n", "Frame 170: Found 8 points to track\n", "Frame 170: Found 8 points to track\n", "Frame 170: Found 8 points to track\n", "Frame 171: Found 8 points to track\n", "Frame 171: Found 8 points to track\n", "Frame 171: Found 8 points to track\n", "Video saved to C:\\Users\\miche\\OneDrive\\Documenti\\GitHub\\bowling-analysis\\data\\recording_3\\test_video\\Spin_test_3.mp4 \n", "Axis log saved to C:\\Users\\miche\\OneDrive\\Documenti\\GitHub\\bowling-analysis\\notebook\\spin\\intermediate_data\\Spin_data_mic_3.csv\n" ] } ], "source": [ "# --- prepare writer ---\n", "fps = cap.get(cv2.CAP_PROP_FPS)\n", "W = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))\n", "H = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))\n", "out_path = str(\n", " project_root\n", " / \"data\"\n", " / f\"recording_{video_number}\"\n", " / \"test_video\"\n", " / f\"Spin_test_{video_number}.mp4\"\n", ")\n", "fourcc = cv2.VideoWriter_fourcc(*\"mp4v\")\n", "out = cv2.VideoWriter(str(out_path), fourcc, fps, (W, H))\n", "\n", "# Prepare list to log axis endpoints\n", "log = []\n", "\n", "# Read frame0 and frame1\n", "ret_curr, curr = cap.read()\n", "ret_next, next_ = cap.read()\n", "frame_idx = 0\n", "\n", "while ret_curr and ret_next:\n", " vis = curr.copy()\n", " pA, pB = (np.nan, np.nan), (np.nan, np.nan)\n", "\n", " # Only compute if we have valid circle data\n", " if (\n", " frame_idx < len(df)\n", " and not df.iloc[frame_idx][[\"x\", \"y\", \"radius\"]].isna().any()\n", " ):\n", " x, y, r = df.iloc[frame_idx][[\"x\", \"y\", \"radius\"]].astype(int)\n", " center = np.array([x, y])\n", " off = 2\n", " x0, x1 = max(x - r - off, 0), min(x + r + off, W)\n", " y0, y1 = max(y - r - off, 0), min(y + r + off, H)\n", "\n", " roi1 = curr[y0:y1, x0:x1]\n", " roi2 = next_[y0:y1, x0:x1]\n", " g1 = cv2.cvtColor(roi1, cv2.COLOR_BGR2GRAY)\n", " g2 = cv2.cvtColor(roi2, cv2.COLOR_BGR2GRAY)\n", " c_roi = np.array([x - x0, y - y0])\n", "\n", " axes = []\n", " for _ in range(3):\n", " # Randomly sample points within the circle\n", " pts = []\n", " while len(pts) < 1000:\n", " xi = np.random.randint(0, g1.shape[1])\n", " yi = np.random.randint(0, g1.shape[0])\n", " if (xi - c_roi[0]) ** 2 + (yi - c_roi[1]) ** 2 <= (r - 2) ** 2:\n", " pts.append([xi, yi])\n", " p0 = np.array(pts, dtype=np.float32).reshape(-1, 1, 2)\n", "\n", " if p0.size == 0:\n", " continue\n", "\n", " p1, st, _ = cv2.calcOpticalFlowPyrLK(\n", " g1,\n", " g2,\n", " p0,\n", " None,\n", " winSize=(15, 15),\n", " maxLevel=2,\n", " criteria=(cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 10, 0.03),\n", " )\n", "\n", " old3d, new3d = [], []\n", " for o, n, s in zip(p0.reshape(-1, 2), p1.reshape(-1, 2), st.reshape(-1)):\n", " if s:\n", " ox, oy = o - c_roi\n", " nx, ny = n - c_roi\n", " oz = np.sqrt(max(r * r - ox * ox - oy * oy, 0))\n", " nz = np.sqrt(max(r * r - nx * nx - ny * ny, 0))\n", " old3d.append([ox, oy, oz])\n", " new3d.append([nx, ny, nz])\n", "\n", " if len(old3d) < 3:\n", " continue\n", "\n", " old3d = np.array(old3d)\n", " new3d = np.array(new3d)\n", "\n", " # Compute Rotation matrix using SVD\n", " Hmat = old3d.T @ new3d\n", " U, S, Vt = np.linalg.svd(Hmat)\n", " R = Vt.T @ U.T\n", " if np.linalg.det(R) < 0:\n", " Vt[-1, :] *= -1\n", " R = Vt.T @ U.T\n", "\n", " theta = np.arccos(np.clip((np.trace(R) - 1) / 2, -1, 1))\n", " if np.sin(theta) != 0:\n", " ax = np.array([R[2, 1] - R[1, 2], R[0, 2] - R[2, 0], R[1, 0] - R[0, 1]])\n", " ax /= 2 * np.sin(theta)\n", " ax /= np.linalg.norm(ax)\n", " axes.append(ax)\n", "\n", " if axes:\n", " avg = np.mean(np.vstack(axes), axis=0)\n", " avg /= np.linalg.norm(avg)\n", "\n", " pA_ball = avg * r\n", " pB_ball = -avg * r\n", "\n", " pA = tuple((center + pA_ball[:2]).astype(int))\n", " pB = tuple((center + pB_ball[:2]).astype(int))\n", "\n", " cv2.line(vis, pA, pB, (0, 255, 255), 2)\n", " cv2.circle(vis, pA, 2, (0, 0, 255), 2)\n", " cv2.circle(vis, pB, 2, (0, 255, 0), 2)\n", "\n", " # log the endpoints for this frame\n", " log.append(\n", " {\"frame\": frame_idx, \"pA_x\": pA[0], \"pA_y\": pA[1], \"pB_x\": pB[0], \"pB_y\": pB[1]}\n", " )\n", "\n", " out.write(vis)\n", "\n", " # advance frames\n", " frame_idx += 1\n", " curr = next_\n", " ret_next, next_ = cap.read()\n", " ret_curr = curr is not None\n", "\n", "# write last frame if exists\n", "if ret_curr:\n", " vis = curr.copy()\n", " log.append(\n", " {\n", " \"frame\": frame_idx,\n", " \"pA_x\": np.nan,\n", " \"pA_y\": np.nan,\n", " \"pB_x\": np.nan,\n", " \"pB_y\": np.nan,\n", " }\n", " )\n", " out.write(vis)\n", "\n", "cap.release()\n", "out.release()\n", "\n", "# Save the log to CSV\n", "log_df = pd.DataFrame(log)\n", "csv_path = str(\n", " project_root\n", " / \"notebook\"\n", " / \"spin\"\n", " / \"intermediate_data\"\n", " / f\"Spin_data_mic_{video_number}.csv\"\n", ")\n", "# log_df.to_csv(csv_path, index=False)\n", "\n", "print(f\"Video saved to {out_path} \\nAxis log saved to {csv_path}\")\n", "\n", "# Add rotation center columns to ball_df\n", "rotation_centers = []\n", "for entry in log:\n", " if (\n", " entry[\"pA_x\"] != np.nan\n", " and entry[\"pA_y\"] != np.nan\n", " and entry[\"pB_x\"] != np.nan\n", " and entry[\"pB_y\"] != np.nan\n", " ):\n", " pA_z = np.sqrt(\n", " max(r * r - (entry[\"pA_x\"] - x) ** 2 - (entry[\"pA_y\"] - y) ** 2, 0)\n", " )\n", " pB_z = np.sqrt(\n", " max(r * r - (entry[\"pB_x\"] - x) ** 2 - (entry[\"pB_y\"] - y) ** 2, 0)\n", " )\n", " if pA_z < 0:\n", " rotation_centers.append(\n", " {\"rotation_center_x\": entry[\"pA_x\"], \"rotation_center_y\": entry[\"pA_y\"]}\n", " )\n", " else:\n", " rotation_centers.append(\n", " {\"rotation_center_x\": entry[\"pB_x\"], \"rotation_center_y\": entry[\"pB_y\"]}\n", " )\n", " else:\n", " rotation_centers.append(\n", " {\"rotation_center_x\": np.nan, \"rotation_center_y\": np.nan}\n", " )\n", "\n", "rotation_centers_df = pd.DataFrame(rotation_centers)\n", "ball_df = pd.concat([df, rotation_centers_df], axis=1)\n", "ball_df.to_csv(csv_path, index=False)" ] } ], "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": 5 }