{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":59156,"databundleVersionId":13874099,"sourceType":"competition"}],"dockerImageVersionId":31090,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"This notebook demonstrates how to visualize the mouse pose tracking data.\n\n## 🎯 Goal\n\nThe primary goal of this code is to create a dynamic animation that shows the movement of a mouse's different body parts over a series of video frames. This helps in understanding the raw tracking data and verifying its quality.\n\n## 📝 Process\n\n1.  **Setup**: We first import the necessary libraries and define the path to the dataset. We specify a `VIDEO_ID` to select a single video for analysis.\n2.  **Data Loading**:\n      * The tracking data (x, y coordinates for each body part at each frame) is loaded from a `.parquet` file.\n      * The corresponding video metadata (like resolution) is loaded from the `train.csv` file.\n3.  **Data Preparation**:\n      * To make the animation process faster, we select a subset of the total frames (`N_FRAMES_TO_ANIMATE`).\n      * We assign a unique color to each body part using the `rainbow` colormap. This makes it easy to distinguish different parts in the animation.\n4.  **Animation**:\n      * We use `matplotlib.animation.FuncAnimation` to generate the animation.\n      * A plot is initialized with the correct dimensions based on the video's resolution.\n      * An `update` function is defined to redraw the positions of all body parts for each new frame.\n      * A legend is added to the plot to show the mapping between colors and body parts.\n5.  **Display**: The final animation is rendered as an interactive HTML object directly within the notebook, complete with play and pause controls.\n\nThis visualization is a crucial first step for exploring the dataset, debugging tracking issues, and developing features for behavior classification models.","metadata":{}},{"cell_type":"code","source":"# Import necessary libraries\nimport pandas as pd\nimport numpy as np\nimport matplotlib.pyplot as plt\nimport matplotlib.cm as cm\nfrom matplotlib.animation import FuncAnimation\nfrom IPython.display import HTML\nfrom pathlib import Path\n\n# ===================================================================\n# 1. Set data paths and target video ID\n# ===================================================================\nVIDEO_ID = \"44566106\"\nLAB = \"AdaptableSnail\"\n\n# --- Automatically find the root path of the dataset ---\nCANDIDATE_BASES = [\"/kaggle/input/MABe-mouse-behavior-detection/\"]\nBASE = None\nfor b in CANDIDATE_BASES:\n    p = Path(b)\n    if (p / \"train_tracking\").exists():\n        BASE = p\n        break\nif BASE is None:\n    raise FileNotFoundError(\"Dataset path not found. Please update CANDIDATE_BASES.\")\n\nTRACKING_FILE = BASE / \"train_tracking\" / LAB / f\"{VIDEO_ID}.parquet\"\nMETA_FILE = BASE / \"train.csv\"\n\nprint(f\"Tracking file: {TRACKING_FILE}\")\nprint(f\"Meta file: {META_FILE}\")\n\n# ===================================================================\n# 2. Load and preprocess data\n# ===================================================================\n# --- Load tracking data (coordinates) ---\ndf = pd.read_parquet(TRACKING_FILE)\n\n# --- Load metadata (video information) ---\nmeta_df = pd.read_csv(META_FILE)\nmeta_df[\"video_id\"] = meta_df[\"video_id\"].astype(str)\nvideo_meta = meta_df[meta_df[\"video_id\"] == VIDEO_ID].iloc[0]\n\n# --- Limit the number of frames for the animation (processing all frames can be very time-consuming) ---\nN_FRAMES_TO_ANIMATE = 500  \ndf_sub = df[df['video_frame'] < N_FRAMES_TO_ANIMATE].copy()\n\n# --- Assign colors to each body part ---\nbodyparts = sorted(df_sub['bodypart'].unique())\n# Use the 'rainbow' colormap to generate colors based on the number of body parts\ncolors = cm.rainbow(np.linspace(0, 1, len(bodyparts)))\n# Create a dictionary to map body part names to colors\npart_to_color = dict(zip(bodyparts, colors))\n# Add color information to each row\ndf_sub['color'] = df_sub['bodypart'].map(part_to_color)\n\nprint(f\"\\nAnimating first {N_FRAMES_TO_ANIMATE} frames...\")\nprint(f\"{len(bodyparts)} body parts will be colored.\")\n\n# ===================================================================\n# 3. Prepare the animation\n# ===================================================================\n# --- Create the plotting area ---\nfig, ax = plt.subplots(figsize=(10, 8))\n\n# --- Set the plot limits to match the video resolution ---\nwidth, height = video_meta['video_width_pix'], video_meta['video_height_pix']\nax.set_xlim(0, width)\nax.set_ylim(height, 0)  # Invert the y-axis to set the origin to the top-left corner\nax.set_title(f\"Mice Pose - video_id: {VIDEO_ID} (Frame: 0)\")\nax.set_xlabel(\"x (pixels)\")\nax.set_ylabel(\"y (pixels)\")\nax.set_aspect('equal', adjustable='box') # Set the aspect ratio to 1:1\nfig.tight_layout()\n\n# --- Initialize the plot (initially empty) ---\n# s=20 sets the size of the points\nscatter = ax.scatter([], [], s=20, alpha=0.8)\n\n# --- Create a legend (to show which color corresponds to which body part) ---\nfor part, color in part_to_color.items():\n    ax.scatter([], [], c=[color], label=part, s=30) # Create the legend using dummy plots\nax.legend(bbox_to_anchor=(1.02, 1), loc='upper left', borderaxespad=0.)\n\n\n# ===================================================================\n# 4. Define the animation update logic\n# ===================================================================\ndef update(frame_num):\n    \"\"\"Function called for each frame\"\"\"\n    # Extract data for the current frame\n    current_frame_data = df_sub[df_sub['video_frame'] == frame_num]\n    \n    if not current_frame_data.empty:\n        # Update the coordinates and colors of the scatter plot\n        scatter.set_offsets(current_frame_data[['x', 'y']].values)\n        scatter.set_color(current_frame_data['color'].values)\n        \n    # Update the title\n    ax.set_title(f\"Mice Pose - video_id: {VIDEO_ID} (Frame: {frame_num})\")\n    \n    # Return the updated plot object\n    return scatter,\n\n# ===================================================================\n# 5. Generate and display the animation\n# ===================================================================\n# --- Create the animation object ---\n# interval=33 corresponds to approximately 30fps (1000ms / 30fps)\n# blit=True is an optimization for faster redrawing\nani = FuncAnimation(fig, update, frames=range(N_FRAMES_TO_ANIMATE), \n                    interval=33, blit=True)\n\n# --- Prevent the static plot from being displayed twice ---\nplt.close(fig)\n\n# --- Display the animation in the Jupyter Notebook ---\n# Using to_jshtml() displays the animation with play/pause controls\nHTML(ani.to_jshtml())\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}