{ "cells": [ { "cell_type": "markdown", "id": "04cc336d", "metadata": {}, "source": [ "# **Spin** calculation\n", "\n", "IDEA:\n", "- take in consideration only the ball for each frame\n", "- draw a lot of points on the ball\n", "- take the motion of these points between consecutive frames (optical flow)\n", "- from this motion we extract the direction of the rotation and so the axis of rotation" ] }, { "cell_type": "code", "execution_count": 1, "id": "abe39cec", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "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_3.csv\n" ] } ], "source": [ "import cv2\n", "import numpy as np\n", "import pandas as pd\n", "from pathlib import Path\n", "\n", "VIDEO_NUMBER = \"3\"\n", "PROJECT_ROOT = Path().resolve().parent.parent\n", "INPUT_VIDEO_PATH = str(\n", " PROJECT_ROOT\n", " / \"data\"\n", " / f\"recording_{VIDEO_NUMBER}\"\n", " / f\"Recording_{VIDEO_NUMBER}.mp4\"\n", ")\n", "INPUT_CSV_PATH = str(\n", " PROJECT_ROOT\n", " / \"notebook\"\n", " / \"ball_detection\"\n", " / \"intermediate_data\"\n", " / f\"Circle_positions_cleaned_{VIDEO_NUMBER}.csv\"\n", ")\n", "\n", "# --- open video & CSV ---\n", "cap = cv2.VideoCapture(INPUT_VIDEO_PATH)\n", "if not cap.isOpened():\n", " raise IOError(\"Could not open video.\")\n", "df = pd.read_csv(INPUT_CSV_PATH)\n", "if df.empty:\n", " raise ValueError(\"CSV is empty.\")\n", "\n", "# --- 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", " 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", " 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", " 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", " if avg[2] != 0:\n", " a2 = avg[:2] / np.linalg.norm(avg[:2])\n", " pA = tuple((center + a2 * r).astype(int))\n", " pB = tuple((center - a2 * r).astype(int))\n", " cv2.line(vis, pA, pB, (255, 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_{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}\")" ] } ], "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 }