{"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":31089,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nfrom tqdm import tqdm","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-04T13:10:41.335186Z","iopub.execute_input":"2025-10-04T13:10:41.335479Z","iopub.status.idle":"2025-10-04T13:10:43.353145Z","shell.execute_reply.started":"2025-10-04T13:10:41.335449Z","shell.execute_reply":"2025-10-04T13:10:43.352192Z"},"_kg_hide-input":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\"\"\"F Beta customized for the data format of the MABe challenge.\"\"\"\n\nimport json\n\nfrom collections import defaultdict\n\nimport pandas as pd\nimport polars as pl\n\n\nclass HostVisibleError(Exception):\n pass\n\n\ndef single_lab_f1(lab_solution: pl.DataFrame, lab_submission: pl.DataFrame, beta: float = 1) -> float:\n label_frames: defaultdict[str, set[int]] = defaultdict(set) # key is video/agent/target/action from solution\n prediction_frames: defaultdict[str, set[int]] = defaultdict(set) # key is video/agent/target/action from submission\n\n for row in lab_solution.to_dicts():\n label_frames[row['label_key']].update(range(row['start_frame'], row['stop_frame']))\n\n for video in lab_solution['video_id'].unique():\n active_labels: str = lab_solution.filter(pl.col('video_id') == video)['behaviors_labeled'].first() # ty: ignore\n active_labels: set[str] = set(json.loads(active_labels)) # set of agent,target,action from solution\n predicted_mouse_pairs: defaultdict[str, set[int]] = defaultdict(set) # key is agent,target from submission\n\n for row in lab_submission.filter(pl.col('video_id') == video).to_dicts(): # every submission row is converted to a dict\n # Since the labels are sparse, we can't evaluate prediction keys not in the active labels.\n if ','.join([str(row['agent_id']), str(row['target_id']), row['action']]) not in active_labels:\n continue # these submission rows are ignored\n \n new_frames = set(range(row['start_frame'], row['stop_frame']))\n # Ignore truly redundant predictions.\n new_frames = new_frames.difference(prediction_frames[row['prediction_key']])\n prediction_pair = ','.join([str(row['agent_id']), str(row['target_id'])])\n if predicted_mouse_pairs[prediction_pair].intersection(new_frames):\n # A single agent can have multiple targets per frame (ex: evading all other mice) but only one action per target per frame.\n raise HostVisibleError('Multiple predictions for the same frame from one agent/target pair')\n prediction_frames[row['prediction_key']].update(new_frames)\n predicted_mouse_pairs[prediction_pair].update(new_frames)\n\n tps = defaultdict(int) # key is action\n fns = defaultdict(int) # key is action\n fps = defaultdict(int) # key is action\n for key, pred_frames in prediction_frames.items():\n action = key.split('_')[-1]\n matched_label_frames = label_frames[key]\n tps[action] += len(pred_frames.intersection(matched_label_frames))\n fns[action] += len(matched_label_frames.difference(pred_frames))\n fps[action] += len(pred_frames.difference(matched_label_frames))\n\n distinct_actions = set()\n for key, frames in label_frames.items():\n action = key.split('_')[-1]\n distinct_actions.add(action)\n if key not in prediction_frames:\n fns[action] += len(frames)\n\n action_f1s = []\n for action in distinct_actions:\n # print(f\"{tps[action]:8} {fns[action]:8} {fps[action]:8}\")\n if tps[action] + fns[action] + fps[action] == 0:\n action_f1s.append(0)\n else:\n action_f1s.append((1 + beta**2) * tps[action] / ((1 + beta**2) * tps[action] + beta**2 * fns[action] + fps[action]))\n return sum(action_f1s) / len(action_f1s)\n\n\ndef mouse_fbeta(solution: pd.DataFrame, submission: pd.DataFrame, beta: float = 1) -> float:\n \"\"\"\n Doctests:\n >>> solution = pd.DataFrame([\n ... {'video_id': 1, 'agent_id': 1, 'target_id': 2, 'action': 'attack', 'start_frame': 0, 'stop_frame': 10, 'lab_id': 1, 'behaviors_labeled': '[\"1,2,attack\"]'},\n ... ])\n >>> submission = pd.DataFrame([\n ... {'video_id': 1, 'agent_id': 1, 'target_id': 2, 'action': 'attack', 'start_frame': 0, 'stop_frame': 10},\n ... ])\n >>> mouse_fbeta(solution, submission)\n 1.0\n\n >>> solution = pd.DataFrame([\n ... {'video_id': 1, 'agent_id': 1, 'target_id': 2, 'action': 'attack', 'start_frame': 0, 'stop_frame': 10, 'lab_id': 1, 'behaviors_labeled': '[\"1,2,attack\"]'},\n ... ])\n >>> submission = pd.DataFrame([\n ... {'video_id': 1, 'agent_id': 1, 'target_id': 2, 'action': 'mount', 'start_frame': 0, 'stop_frame': 10}, # Wrong action\n ... ])\n >>> mouse_fbeta(solution, submission)\n 0.0\n\n >>> solution = pd.DataFrame([\n ... {'video_id': 123, 'agent_id': 1, 'target_id': 2, 'action': 'attack', 'start_frame': 0, 'stop_frame': 9, 'lab_id': 1, 'behaviors_labeled': '[\"1,2,attack\"]'},\n ... {'video_id': 123, 'agent_id': 1, 'target_id': 2, 'action': 'mount', 'start_frame': 15, 'stop_frame': 24, 'lab_id': 1, 'behaviors_labeled': '[\"1,2,attack\"]'},\n ... ])\n >>> submission = pd.DataFrame([\n ... {'video_id': 123, 'agent_id': 1, 'target_id': 2, 'action': 'attack', 'start_frame': 0, 'stop_frame': 9},\n ... ])\n >>> \"%.12f\" % mouse_fbeta(solution, submission)\n '0.500000000000'\n\n >>> solution = pd.DataFrame([\n ... {'video_id': 123, 'agent_id': 1, 'target_id': 2, 'action': 'attack', 'start_frame': 0, 'stop_frame': 9, 'lab_id': 1, 'behaviors_labeled': '[\"1,2,attack\"]'},\n ... {'video_id': 123, 'agent_id': 1, 'target_id': 2, 'action': 'mount', 'start_frame': 15, 'stop_frame': 24, 'lab_id': 1, 'behaviors_labeled': '[\"1,2,attack\"]'},\n ... {'video_id': 345, 'agent_id': 1, 'target_id': 2, 'action': 'attack', 'start_frame': 0, 'stop_frame': 9, 'lab_id': 2, 'behaviors_labeled': '[\"1,2,attack\"]'},\n ... {'video_id': 345, 'agent_id': 1, 'target_id': 2, 'action': 'mount', 'start_frame': 15, 'stop_frame': 24, 'lab_id': 2, 'behaviors_labeled': '[\"1,2,attack\"]'},\n ... ])\n >>> submission = pd.DataFrame([\n ... {'video_id': 123, 'agent_id': 1, 'target_id': 2, 'action': 'attack', 'start_frame': 0, 'stop_frame': 9},\n ... ])\n >>> \"%.12f\" % mouse_fbeta(solution, submission)\n '0.250000000000'\n\n >>> # Overlapping solution events, one prediction matching both.\n >>> solution = pd.DataFrame([\n ... {'video_id': 1, 'agent_id': 1, 'target_id': 2, 'action': 'attack', 'start_frame': 0, 'stop_frame': 10, 'lab_id': 1, 'behaviors_labeled': '[\"1,2,attack\"]'},\n ... {'video_id': 1, 'agent_id': 1, 'target_id': 2, 'action': 'attack', 'start_frame': 10, 'stop_frame': 20, 'lab_id': 1, 'behaviors_labeled': '[\"1,2,attack\"]'},\n ... ])\n >>> submission = pd.DataFrame([\n ... {'video_id': 1, 'agent_id': 1, 'target_id': 2, 'action': 'attack', 'start_frame': 0, 'stop_frame': 20},\n ... ])\n >>> mouse_fbeta(solution, submission)\n 1.0\n\n >>> solution = pd.DataFrame([\n ... {'video_id': 1, 'agent_id': 1, 'target_id': 2, 'action': 'attack', 'start_frame': 0, 'stop_frame': 10, 'lab_id': 1, 'behaviors_labeled': '[\"1,2,attack\"]'},\n ... {'video_id': 1, 'agent_id': 1, 'target_id': 2, 'action': 'attack', 'start_frame': 30, 'stop_frame': 40, 'lab_id': 1, 'behaviors_labeled': '[\"1,2,attack\"]'},\n ... ])\n >>> submission = pd.DataFrame([\n ... {'video_id': 1, 'agent_id': 1, 'target_id': 2, 'action': 'attack', 'start_frame': 0, 'stop_frame': 40},\n ... ])\n >>> mouse_fbeta(solution, submission)\n 0.6666666666666666\n \"\"\"\n if len(solution) == 0 or len(submission) == 0:\n raise ValueError('Missing solution or submission data')\n\n expected_cols = ['video_id', 'agent_id', 'target_id', 'action', 'start_frame', 'stop_frame']\n\n for col in expected_cols:\n if col not in solution.columns:\n raise ValueError(f'Solution is missing column {col}')\n if col not in submission.columns:\n raise ValueError(f'Submission is missing column {col}')\n\n solution: pl.DataFrame = pl.DataFrame(solution)\n submission: pl.DataFrame = pl.DataFrame(submission)\n assert (solution['start_frame'] <= solution['stop_frame']).all()\n assert (submission['start_frame'] <= submission['stop_frame']).all()\n solution_videos = set(solution['video_id'].unique())\n # Need to align based on video IDs as we can't rely on the row IDs for handling public/private splits.\n submission = submission.filter(pl.col('video_id').is_in(solution_videos))\n\n solution = solution.with_columns(\n pl.concat_str(\n [\n pl.col('video_id').cast(pl.Utf8),\n pl.col('agent_id').cast(pl.Utf8),\n pl.col('target_id').cast(pl.Utf8),\n pl.col('action'),\n ],\n separator='_',\n ).alias('label_key'),\n )\n submission = submission.with_columns(\n pl.concat_str(\n [\n pl.col('video_id').cast(pl.Utf8),\n pl.col('agent_id').cast(pl.Utf8),\n pl.col('target_id').cast(pl.Utf8),\n pl.col('action'),\n ],\n separator='_',\n ).alias('prediction_key'),\n )\n\n lab_scores = []\n for lab in solution['lab_id'].unique():\n lab_solution = solution.filter(pl.col('lab_id') == lab).clone()\n lab_videos = set(lab_solution['video_id'].unique())\n lab_submission = submission.filter(pl.col('video_id').is_in(lab_videos)).clone()\n # print(len(lab_solution), len(lab_videos), len(lab_submission), single_lab_f1(lab_solution, lab_submission, beta=beta))\n lab_scores.append(single_lab_f1(lab_solution, lab_submission, beta=beta))\n\n return sum(lab_scores) / len(lab_scores)\n\n\ndef score(solution: pd.DataFrame, submission: pd.DataFrame, row_id_column_name: str, beta: float = 1) -> float:\n \"\"\"\n F1 score for the MABe Challenge\n \"\"\"\n solution = solution.drop(row_id_column_name, axis='columns', errors='ignore')\n submission = submission.drop(row_id_column_name, axis='columns', errors='ignore')\n return mouse_fbeta(solution, submission, beta=beta)","metadata":{"trusted":true,"_kg_hide-input":true,"execution":{"iopub.status.busy":"2025-10-04T13:12:06.031626Z","iopub.execute_input":"2025-10-04T13:12:06.032332Z","iopub.status.idle":"2025-10-04T13:12:06.724393Z","shell.execute_reply.started":"2025-10-04T13:12:06.032300Z","shell.execute_reply":"2025-10-04T13:12:06.723499Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Validation\n\nWe first create the solution dataframe, which contains all true labels in the correct format for the `score()` function.","metadata":{}},{"cell_type":"code","source":"train = pd.read_csv('/kaggle/input/MABe-mouse-behavior-detection/train.csv')\ntrain_subset = train.query(\"~ lab_id.str.startswith('MABe22_')\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-04T13:12:12.539801Z","iopub.execute_input":"2025-10-04T13:12:12.540194Z","iopub.status.idle":"2025-10-04T13:12:12.701489Z","shell.execute_reply.started":"2025-10-04T13:12:12.540162Z","shell.execute_reply":"2025-10-04T13:12:12.700507Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def create_solution_df(dataset):\n \"\"\"Create the solution dataframe for validating out-of-fold predictions.\n\n From https://www.kaggle.com/code/ambrosm/mabe-validated-baseline-without-machine-learning/\n \n Parameters:\n dataset: (a subset of) the train dataframe\n \n Return values:\n solution: solution dataframe in the correct format for the score() function\n \"\"\"\n solution = []\n for train_idx, row in tqdm(dataset.iterrows(), total=len(dataset)):\n \n # Load annotation file\n lab_id = row['lab_id']\n if lab_id.startswith('MABe22'): continue\n video_id = row['video_id']\n path = f\"/kaggle/input/MABe-mouse-behavior-detection/train_annotation/{lab_id}/{video_id}.parquet\"\n try:\n annot = pd.read_parquet(path.replace('train_tracking', 'train_annotation'))\n except FileNotFoundError:\n print(f\"No annotations for {path}\")\n continue\n \n # Add all annotations to the solution\n self_annotation = annot.target_id == annot.agent_id\n annot['lab_id'] = lab_id\n annot['video_id'] = video_id\n annot['behaviors_labeled'] = row['behaviors_labeled']\n annot['agent_id'] = annot['agent_id'].apply(lambda s: f\"mouse{s}\")\n annot['target_id'] = np.where(self_annotation, 'self', annot['target_id'].apply(lambda s: f\"mouse{s}\"))\n solution.append(annot)\n \n solution = pd.concat(solution)\n return solution\n\nsolution = create_solution_df(train_subset)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-04T13:12:14.915728Z","iopub.execute_input":"2025-10-04T13:12:14.916117Z","iopub.status.idle":"2025-10-04T13:12:22.569116Z","shell.execute_reply.started":"2025-10-04T13:12:14.916092Z","shell.execute_reply":"2025-10-04T13:12:22.568247Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Now we'll predict something. In binary classification without machine learning, you can either predict all true or all false:\n- If you predict all false, the f1 score will be 0 because the recall is 0.\n- If you predict all true, the recall will be 1 and the precision > 0 so that the f1 score (harmonic mean of recall and precision) will be positive.\n\nPredicting all true is obviously the better choice.\n\nHowever, this competition is not binary classification. For every video_id/agent/target combination, we may predict *at most one* of several actions. This is multi-class classification (every action is a class, and 'no-action' is an additional class). To get positive recalls and f1 scores for all actions, we distribute our predictions to all of them (but we never predict 'no-action'). Note that the set of possible actions depends on the video_id/agent/target combination.\n","metadata":{}},{"cell_type":"code","source":"def predict_without_ml(dataset, traintest):\n \"\"\"Predict actions without machine learning\n \n Parameters:\n dataset: (a subset of) the train or test dataframe\n traintest: 'train' or 'test'\n \n Return values:\n submission: submission dataframe\n \"\"\"\n submission = []\n for train_idx, row in tqdm(dataset.iterrows(), total=len(dataset)):\n \n # Load video\n lab_id = row['lab_id']\n if lab_id.startswith('MABe22'): continue\n video_id = row['video_id']\n path = f\"/kaggle/input/MABe-mouse-behavior-detection/{traintest}_tracking/{lab_id}/{video_id}.parquet\"\n vid = pd.read_parquet(path)\n \n # Determine the behaviors of this video\n if type(row.behaviors_labeled) != str:\n continue\n vid_behaviors = json.loads(row['behaviors_labeled'])\n vid_behaviors = sorted(list({b.replace(\"'\", \"\") for b in vid_behaviors}))\n vid_behaviors = [b.split(',') for b in vid_behaviors]\n vid_behaviors = pd.DataFrame(vid_behaviors, columns=['agent', 'target', 'action'])\n \n # Determine start_frame and stop_frame\n start_frame = vid.video_frame.min()\n stop_frame = vid.video_frame.max() + 1\n \n # Predict all possible actions as often as possible\n for (agent, target), actions in vid_behaviors.groupby(['agent', 'target']):\n batch_length = int(np.ceil((stop_frame - start_frame) / len(actions)))\n for i, (_, action_row) in enumerate(actions.iterrows()):\n batch_start = start_frame + i * batch_length\n batch_stop = min(batch_start + batch_length, stop_frame)\n submission.append((video_id, agent, target, action_row['action'], batch_start, batch_stop))\n \n # Convert to dataframe\n submission = pd.DataFrame(submission, columns=['video_id', 'agent_id', 'target_id', 'action', 'start_frame', 'stop_frame'])\n\n return submission\n\nsubmission = predict_without_ml(train_subset, 'train')\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-04T13:16:22.522635Z","iopub.execute_input":"2025-10-04T13:16:22.523006Z","iopub.status.idle":"2025-10-04T13:16:44.533220Z","shell.execute_reply.started":"2025-10-04T13:16:22.522978Z","shell.execute_reply":"2025-10-04T13:16:44.532284Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"What score would we get on the training dataset?","metadata":{}},{"cell_type":"code","source":"%%time\nscore(solution, submission, '')\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-04T13:16:50.151386Z","iopub.execute_input":"2025-10-04T13:16:50.151729Z","iopub.status.idle":"2025-10-04T13:17:01.329073Z","shell.execute_reply.started":"2025-10-04T13:16:50.151702Z","shell.execute_reply":"2025-10-04T13:17:01.328052Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Submission","metadata":{}},{"cell_type":"code","source":"test = pd.read_csv('/kaggle/input/MABe-mouse-behavior-detection/test.csv')\nsubmission = predict_without_ml(test, 'test')\nsubmission.index.name = 'row_id'\nsubmission.to_csv('submission.csv')\n!head submission.csv","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-04T13:17:41.203613Z","iopub.execute_input":"2025-10-04T13:17:41.205990Z","iopub.status.idle":"2025-10-04T13:17:41.492094Z","shell.execute_reply.started":"2025-10-04T13:17:41.205941Z","shell.execute_reply":"2025-10-04T13:17:41.490986Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from pathlib import Path\nimport pandas as pd\nimport numpy as np\nimport matplotlib.pyplot as plt\nimport seaborn as sns\n\n# consistent color palette\nPALETTE = {\n \"deep_blue\": \"#1E3A8A\",\n \"slate_gray\": \"#374151\",\n \"teal\": \"#0D9488\",\n \"muted\": \"#94A3B8\",\n \"bg\": \"#FFFFFF\"\n}\n\n# seaborn/matplotlib defaults\nsns.set_theme(\n context=\"talk\",\n style=\"whitegrid\",\n rc={\n \"figure.dpi\": 120,\n \"axes.facecolor\": PALETTE[\"bg\"],\n \"figure.facecolor\": PALETTE[\"bg\"],\n \"font.family\": \"sans-serif\",\n \"font.size\": 12,\n \"axes.titlesize\": 16,\n \"axes.labelsize\": 13,\n \"legend.fontsize\": 11,\n \"xtick.labelsize\": 11,\n \"ytick.labelsize\": 11,\n \"grid.color\": \"#E6EEF2\",\n \"grid.linewidth\": 0.8\n }\n)\n\n# dataset root\nDATA = Path(\"/kaggle/input/MABe-mouse-behavior-detection\")\npd.set_option(\"display.max_columns\", None)\npd.set_option(\"display.width\", 240)\n\n# read train.csv\ntrain_csv_path = DATA / \"train.csv\"\nif not train_csv_path.exists():\n raise FileNotFoundError(f\"train.csv not found at: {train_csv_path}\")\ntrain_csv = pd.read_csv(train_csv_path)\ntrain_csv[\"video_id\"] = train_csv[\"video_id\"].astype(str)\n\n# folder structure\nfolders = {\n \"train_tracking\": DATA / \"train_tracking\",\n \"train_annotation\": DATA / \"train_annotation\",\n \"test_tracking\": DATA / \"test_tracking\"\n}\n\n# collect summary\nsummary_data = []\nfor folder_name, folder_path in folders.items():\n if not folder_path.exists():\n summary_data.append((folder_name, 0, 0))\n continue\n total_files = 0\n lab_dirs = [p for p in folder_path.iterdir() if p.is_dir()]\n for lab_dir in lab_dirs:\n files_in_lab = len(list(lab_dir.glob(\"*.parquet\")))\n total_files += files_in_lab\n summary_data.append((folder_name, total_files, len(lab_dirs)))\n\nsummary_df = pd.DataFrame(summary_data, columns=[\"Directory\", \"Total Files\", \"Lab Count\"]).sort_values(\"Total Files\", ascending=True)\n\n# plot\nfig, ax = plt.subplots(figsize=(8, 4.5))\nsns.barplot(data=summary_df, x=\"Total Files\", y=\"Directory\", palette=[PALETTE[\"deep_blue\"]], ax=ax)\n\nfor p in ax.patches:\n ax.annotate(f\"{int(p.get_width()):,}\", \n (p.get_width(), p.get_y() + p.get_height() / 2),\n xytext=(6, 0), textcoords=\"offset points\", \n va=\"center\", color=PALETTE[\"slate_gray\"], fontsize=11)\n\nax.set_title(\"Top-level file counts by directory\", color=PALETTE[\"slate_gray\"], fontsize=14)\nax.set_xlabel(\"Total parquet files\")\nax.set_ylabel(\"\")\nax.spines[\"top\"].set_visible(False)\nax.spines[\"right\"].set_visible(False)\nax.grid(axis=\"x\", linestyle=\"--\", alpha=0.35)\n\nplt.tight_layout()\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-04T13:17:46.229867Z","iopub.execute_input":"2025-10-04T13:17:46.230243Z","iopub.status.idle":"2025-10-04T13:17:48.195432Z","shell.execute_reply.started":"2025-10-04T13:17:46.230197Z","shell.execute_reply":"2025-10-04T13:17:48.194498Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"lab_inventory = []\nfor lab_id in sorted(train_csv[\"lab_id\"].unique()):\n tracking_dir = folders[\"train_tracking\"] / lab_id\n annotation_dir = folders[\"train_annotation\"] / lab_id\n tracking_count = len(list(tracking_dir.glob(\"*.parquet\"))) if tracking_dir.exists() else 0\n annotation_count = len(list(annotation_dir.glob(\"*.parquet\"))) if annotation_dir.exists() else 0\n difference = tracking_count - annotation_count\n lab_inventory.append((lab_id, tracking_count, annotation_count, difference))\n\ninventory_df = pd.DataFrame(\n lab_inventory, \n columns=[\"Lab ID\", \"Tracking Files\", \"Annotation Files\", \"Difference (T-A)\"]\n).sort_values([\"Annotation Files\", \"Tracking Files\"], ascending=[False, False]).reset_index(drop=True)\n\nfig, ax = plt.subplots(figsize=(10, 6))\nsns.barplot(\n data=inventory_df, \n x=\"Lab ID\", \n y=\"Tracking Files\", \n color=PALETTE[\"deep_blue\"], \n ax=ax, \n label=\"Tracking Files\"\n)\nsns.barplot(\n data=inventory_df, \n x=\"Lab ID\", \n y=\"Annotation Files\", \n color=PALETTE[\"teal\"], \n ax=ax, \n label=\"Annotation Files\"\n)\nax.set_title(\"Per-lab tracking vs annotation files\", color=PALETTE[\"slate_gray\"], fontsize=15)\nax.set_xlabel(\"Lab ID\")\nax.set_ylabel(\"Number of files\")\nax.legend(frameon=False)\nax.spines[\"top\"].set_visible(False)\nax.spines[\"right\"].set_visible(False)\nax.grid(axis=\"y\", linestyle=\"--\", alpha=0.35)\nplt.xticks(rotation=45)\nplt.tight_layout()\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-04T13:17:58.919836Z","iopub.execute_input":"2025-10-04T13:17:58.920221Z","iopub.status.idle":"2025-10-04T13:17:59.703806Z","shell.execute_reply.started":"2025-10-04T13:17:58.920184Z","shell.execute_reply":"2025-10-04T13:17:59.702719Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"viz_df = inventory_df.copy()\nviz_df[\"Lab Short\"] = viz_df[\"Lab ID\"].str.replace(\"_\", \" \").str.slice(0, 30)\ntop_labs = viz_df.nlargest(25, \"Tracking Files\").reset_index(drop=True)\n\nfig, ax = plt.subplots(figsize=(12, 10), dpi=120)\ny = np.arange(len(top_labs))\nbar_h = 0.35\n\nax.barh(y - bar_h/2, top_labs[\"Tracking Files\"], height=bar_h, label=\"Tracking Files\", color=PALETTE[\"deep_blue\"], alpha=0.85)\nax.barh(y + bar_h/2, top_labs[\"Annotation Files\"], height=bar_h, label=\"Annotation Files\", color=PALETTE[\"teal\"], alpha=0.95)\n\nfor i, (t, a) in enumerate(zip(top_labs[\"Tracking Files\"], top_labs[\"Annotation Files\"])):\n ax.text(t + max(top_labs[\"Tracking Files\"]) * 0.003, i - bar_h/2, f\"{int(t):,}\", va=\"center\", ha=\"left\", color=PALETTE[\"slate_gray\"], fontsize=10)\n ax.text(a + max(top_labs[\"Tracking Files\"]) * 0.003, i + bar_h/2, f\"{int(a):,}\", va=\"center\", ha=\"left\", color=PALETTE[\"slate_gray\"], fontsize=10)\n\nax.set_yticks(y)\nax.set_yticklabels(top_labs[\"Lab Short\"], color=PALETTE[\"slate_gray\"])\nax.invert_yaxis()\nax.set_xlabel(\"Number of Parquet Files\", color=PALETTE[\"slate_gray\"])\nax.set_title(\"File Availability by Lab — Tracking vs Annotation\", color=PALETTE[\"slate_gray\"], fontsize=16, pad=12)\nax.legend(frameon=False, loc=\"lower right\")\nplt.tight_layout()\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-04T13:18:05.238812Z","iopub.execute_input":"2025-10-04T13:18:05.239630Z","iopub.status.idle":"2025-10-04T13:18:05.865314Z","shell.execute_reply.started":"2025-10-04T13:18:05.239597Z","shell.execute_reply":"2025-10-04T13:18:05.864212Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from IPython.display import HTML, display\n\nanno_root = folders[\"train_annotation\"]\nvalid_annos = set()\nif anno_root.exists():\n for lab_dir in sorted(anno_root.iterdir()):\n if not lab_dir.is_dir():\n continue\n for f in lab_dir.glob(\"*.parquet\"):\n valid_annos.add((lab_dir.name, f.stem))\n\nlabs_no_annotations = inventory_df.loc[inventory_df[\"Annotation Files\"] == 0, \"Lab ID\"].tolist()\nusable_mask = train_csv.apply(lambda r: (r[\"lab_id\"], r[\"video_id\"]) in valid_annos, axis=1)\nusable_train_rows = train_csv[usable_mask]\n\ntotal_rows = len(train_csv)\nusable_rows = len(usable_train_rows)\nusable_pct = (usable_rows / total_rows * 100) if total_rows > 0 else 0.0\n\nsummary_html = f\"\"\"\n
\n

Missing Annotation Analysis

\n
\n
Labs with zero annotation files: {', '.join(labs_no_annotations) if labs_no_annotations else 'None'}
\n
Total rows in train.csv: {total_rows:,}
\n
Usable rows with annotation parquet: {usable_rows:,}
\n
Usable percentage: {usable_pct:.1f}%
\n
\n
\n\"\"\"\ndisplay(HTML(summary_html))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-04T13:18:11.013510Z","iopub.execute_input":"2025-10-04T13:18:11.013855Z","iopub.status.idle":"2025-10-04T13:18:11.115185Z","shell.execute_reply.started":"2025-10-04T13:18:11.013831Z","shell.execute_reply":"2025-10-04T13:18:11.114398Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"top_labs = inventory_df.head(15)\nbottom_labs = inventory_df.tail(15)\n\nfig, axes = plt.subplots(1, 2, figsize=(14, 6), dpi=120, sharey=True)\n\nsns.barplot(data=top_labs, x=\"Tracking Files\", y=\"Lab ID\", color=PALETTE[\"deep_blue\"], ax=axes[0], label=\"Tracking\")\nsns.barplot(data=top_labs, x=\"Annotation Files\", y=\"Lab ID\", color=PALETTE[\"teal\"], ax=axes[0], label=\"Annotation\")\naxes[0].set_title(\"Top 15 Labs by File Count\", color=PALETTE[\"slate_gray\"], fontsize=14)\naxes[0].set_xlabel(\"Number of Files\")\naxes[0].set_ylabel(\"\")\naxes[0].legend(frameon=False)\n\nsns.barplot(data=bottom_labs, x=\"Tracking Files\", y=\"Lab ID\", color=PALETTE[\"deep_blue\"], ax=axes[1], label=\"Tracking\")\nsns.barplot(data=bottom_labs, x=\"Annotation Files\", y=\"Lab ID\", color=PALETTE[\"teal\"], ax=axes[1], label=\"Annotation\")\naxes[1].set_title(\"Bottom 15 Labs by File Count\", color=PALETTE[\"slate_gray\"], fontsize=14)\naxes[1].set_xlabel(\"Number of Files\")\naxes[1].set_ylabel(\"\")\naxes[1].legend(frameon=False)\n\nfor ax in axes:\n ax.spines[\"top\"].set_visible(False)\n ax.spines[\"right\"].set_visible(False)\n ax.grid(axis=\"x\", linestyle=\"--\", alpha=0.35)\n\nplt.tight_layout()\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-04T13:18:15.661925Z","iopub.execute_input":"2025-10-04T13:18:15.662601Z","iopub.status.idle":"2025-10-04T13:18:16.522948Z","shell.execute_reply.started":"2025-10-04T13:18:15.662575Z","shell.execute_reply":"2025-10-04T13:18:16.521888Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"all_annotations = []\nannotation_stats = []\n\nfor _, row in usable_train_rows.iterrows():\n lab_id = row[\"lab_id\"]\n video_id = row[\"video_id\"]\n anno_path = DATA / \"train_annotation\" / lab_id / f\"{video_id}.parquet\"\n try:\n df_anno = pd.read_parquet(anno_path)\n df_anno[\"lab_id\"] = lab_id\n df_anno[\"video_id\"] = video_id\n df_anno[\"duration\"] = df_anno[\"stop_frame\"] - df_anno[\"start_frame\"] + 1\n df_anno[\"is_self_directed\"] = df_anno[\"agent_id\"] == df_anno[\"target_id\"]\n all_annotations.append(df_anno)\n annotation_stats.append((lab_id, video_id, len(df_anno)))\n except Exception:\n continue\n\ncombined_annotations = pd.concat(all_annotations, ignore_index=True)\nstats_df = pd.DataFrame(annotation_stats, columns=[\"lab_id\", \"video_id\", \"annotation_count\"])\n\nannotation_overview = {\n \"Behavioral Instances\": len(combined_annotations),\n \"Unique Actions\": combined_annotations[\"action\"].nunique(),\n \"Videos with Annotations\": len(stats_df),\n \"Labs with Annotations\": combined_annotations[\"lab_id\"].nunique()\n}\n\nfig, ax = plt.subplots(figsize=(7, 7), dpi=120)\nsizes = list(annotation_overview.values())\nlabels = [f\"{k}\\n{v:,}\" for k, v in annotation_overview.items()]\ncolors = [PALETTE[\"deep_blue\"], PALETTE[\"teal\"], \"#60A5FA\", \"#F59E0B\"]\n\nwedges, _ = ax.pie(\n sizes,\n startangle=90,\n counterclock=False,\n wedgeprops=dict(width=0.45),\n colors=colors\n)\nax.set(aspect=\"equal\")\nax.set_title(\"Behavioral Annotation Overview\", color=PALETTE[\"slate_gray\"], fontsize=15, pad=15)\nax.legend(wedges, labels, loc=\"center left\", bbox_to_anchor=(1, 0.5), frameon=False)\nplt.tight_layout()\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-04T13:18:21.429643Z","iopub.execute_input":"2025-10-04T13:18:21.429960Z","iopub.status.idle":"2025-10-04T13:18:26.324244Z","shell.execute_reply.started":"2025-10-04T13:18:21.429938Z","shell.execute_reply":"2025-10-04T13:18:26.323335Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"action_counts = combined_annotations[\"action\"].value_counts()\naction_stats = (\n combined_annotations.groupby(\"action\")[\"duration\"]\n .agg([\"count\", \"mean\", \"std\", \"min\", \"max\"])\n .round(2)\n)\naction_stats[\"percentage\"] = (action_stats[\"count\"] / len(combined_annotations) * 100).round(2)\naction_stats = action_stats.sort_values(\"count\", ascending=False).reset_index()\n\ntop20 = action_stats.head(20)\n\nfig, ax = plt.subplots(figsize=(10, 7), dpi=120)\nsns.barplot(\n data=top20,\n x=\"count\",\n y=\"action\",\n palette=sns.color_palette([PALETTE[\"deep_blue\"], PALETTE[\"teal\"]]),\n ax=ax\n)\nfor i, (c, pct) in enumerate(zip(top20[\"count\"], top20[\"percentage\"])):\n ax.text(c + max(top20[\"count\"]) * 0.003, i, f\"{c:,} ({pct:.1f}%)\",\n va=\"center\", ha=\"left\", fontsize=10, color=PALETTE[\"slate_gray\"])\n\nax.set_title(\"Top 20 Most Frequent Actions\", color=PALETTE[\"slate_gray\"], fontsize=15, pad=12)\nax.set_xlabel(\"Number of Instances\")\nax.set_ylabel(\"\")\nax.spines[\"top\"].set_visible(False)\nax.spines[\"right\"].set_visible(False)\nax.grid(axis=\"x\", linestyle=\"--\", alpha=0.35)\nplt.tight_layout()\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-04T13:18:31.333405Z","iopub.execute_input":"2025-10-04T13:18:31.333786Z","iopub.status.idle":"2025-10-04T13:18:31.863630Z","shell.execute_reply.started":"2025-10-04T13:18:31.333760Z","shell.execute_reply":"2025-10-04T13:18:31.862592Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import matplotlib.pyplot as plt\nimport numpy as np\nfrom IPython.display import HTML, display\n\nif \"combined_annotations\" not in globals():\n if \"usable_train_rows\" in globals():\n all_annotations = []\n for _, row in usable_train_rows.iterrows():\n lab_id = row[\"lab_id\"]\n video_id = row[\"video_id\"]\n anno_path = DATA / \"train_annotation\" / lab_id / f\"{video_id}.parquet\"\n try:\n df_anno = pd.read_parquet(anno_path)\n df_anno[\"lab_id\"] = lab_id\n df_anno[\"video_id\"] = video_id\n df_anno[\"duration\"] = df_anno[\"stop_frame\"] - df_anno[\"start_frame\"] + 1\n df_anno[\"is_self_directed\"] = df_anno[\"agent_id\"] == df_anno[\"target_id\"]\n all_annotations.append(df_anno)\n except Exception:\n continue\n if len(all_annotations) == 0:\n raise NameError(\"No annotations found when attempting to rebuild combined_annotations.\")\n combined_annotations = pd.concat(all_annotations, ignore_index=True)\n else:\n raise NameError(\"combined_annotations not found in the session. Re-run the annotation-loading cell first.\")\n\nsd_counts = combined_annotations[\"is_self_directed\"].value_counts()\nsocial = int(sd_counts.get(False, 0))\nself_dir = int(sd_counts.get(True, 0))\ntotal = social + self_dir\nsocial_pct = social / total * 100 if total > 0 else 0.0\nself_pct = self_dir / total * 100 if total > 0 else 0.0\n\nsummary_html = f\"\"\"\n
\n

Self-directed vs Social Behaviors

\n
\n
Social behaviors (agent != target): {social:,} ({social_pct:.1f}%)
\n
Self-directed behaviors (agent == target): {self_dir:,} ({self_pct:.1f}%)
\n
Total behavioral instances: {total:,}
\n
\n
\n\"\"\"\ndisplay(HTML(summary_html))\n\nfig, ax = plt.subplots(figsize=(6, 6), dpi=120)\nsizes = [social, self_dir]\nlabels = [f\"Social\\n{social:,}\\n{social_pct:.1f}%\", f\"Self-directed\\n{self_dir:,}\\n{self_pct:.1f}%\"]\ncolors = [PALETTE[\"deep_blue\"], PALETTE[\"teal\"]]\n\nwedges, texts = ax.pie(sizes, startangle=90, counterclock=False, wedgeprops=dict(width=0.45), colors=colors)\nax.set(aspect=\"equal\")\nax.set_title(\"Social vs Self-directed Behavioral Instances\", color=PALETTE[\"slate_gray\"], fontsize=14, pad=12)\nax.legend(wedges, labels, title=\"\", loc=\"center left\", bbox_to_anchor=(1, 0.5), frameon=False)\nplt.tight_layout()\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-04T13:19:03.821516Z","iopub.execute_input":"2025-10-04T13:19:03.821878Z","iopub.status.idle":"2025-10-04T13:19:04.006975Z","shell.execute_reply.started":"2025-10-04T13:19:03.821854Z","shell.execute_reply":"2025-10-04T13:19:04.005984Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Compute duration statistics per action\nduration_stats = (\n combined_annotations.groupby(\"action\")[\"duration\"]\n .agg([\"mean\", \"median\", \"std\"])\n .round(1)\n .sort_values(\"mean\", ascending=False)\n)\n\n# Prepare data for top 15 actions\ntop_actions = duration_stats.head(15).index.tolist()\nplot_df = combined_annotations[combined_annotations[\"action\"].isin(top_actions)].copy()\nplot_df = plot_df[plot_df[\"duration\"] > 0]\n\n# Plot distributions\nplt.figure(figsize=(12, 8), dpi=120)\nax = sns.boxplot(\n data=plot_df,\n x=\"duration\",\n y=\"action\",\n order=top_actions,\n showfliers=False,\n width=0.6,\n palette=[PALETTE[\"teal\"] if i % 2 == 0 else PALETTE[\"deep_blue\"] for i in range(len(top_actions))]\n)\nax.set_xscale(\"log\")\nax.set_xlabel(\"Duration (frames) — log scale\", color=PALETTE[\"slate_gray\"])\nax.set_ylabel(\"\")\nax.set_title(\"Action Duration Distributions — Top 15 by Mean Duration\", color=PALETTE[\"slate_gray\"], fontsize=15, pad=12)\n\n# annotate medians\nmedians = plot_df.groupby(\"action\")[\"duration\"].median().reindex(top_actions)\nx_off = plot_df[\"duration\"].median() * 0.03\nfor i, (act, med) in enumerate(medians.items()):\n ax.text(med + x_off, i, f\"{int(med):,}\", va=\"center\", ha=\"left\", color=PALETTE[\"slate_gray\"], fontsize=9)\n\nplt.tight_layout()\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-04T13:21:25.446922Z","iopub.execute_input":"2025-10-04T13:21:25.447745Z","iopub.status.idle":"2025-10-04T13:21:26.327466Z","shell.execute_reply.started":"2025-10-04T13:21:25.447714Z","shell.execute_reply":"2025-10-04T13:21:26.326597Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Build lab-action availability matrix\nlab_action_matrix = combined_annotations.groupby([\"lab_id\", \"action\"]).size().unstack(fill_value=0)\nlab_action_presence = (lab_action_matrix > 0).astype(int)\n\ntop_actions = combined_annotations[\"action\"].value_counts().head(30).index.tolist()\nsubset = lab_action_presence[top_actions]\n\nplt.figure(figsize=(16, 10), dpi=120)\nax = sns.heatmap(\n subset,\n cmap=sns.color_palette([PALETTE[\"muted\"], PALETTE[\"teal\"]]),\n linewidths=0.3,\n linecolor=\"#E5E7EB\",\n cbar_kws={\"label\": \"Action Presence\"}\n)\n\nax.set_title(\"Cross-Lab Action Availability Heatmap (Top 30 Actions)\", color=PALETTE[\"slate_gray\"], fontsize=15, pad=16)\nax.set_xlabel(\"Actions\", color=PALETTE[\"slate_gray\"])\nax.set_ylabel(\"Labs\", color=PALETTE[\"slate_gray\"])\nax.tick_params(axis=\"x\", rotation=65)\nplt.tight_layout()\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-04T13:21:30.510827Z","iopub.execute_input":"2025-10-04T13:21:30.511910Z","iopub.status.idle":"2025-10-04T13:21:31.476862Z","shell.execute_reply.started":"2025-10-04T13:21:30.511875Z","shell.execute_reply":"2025-10-04T13:21:31.475574Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class_imbalance = action_counts / action_counts.sum()\nrare_actions = action_counts[action_counts <= 10]\ncommon_actions = action_counts[action_counts >= 1000]\n\nrare_count = len(rare_actions)\ncommon_count = len(common_actions)\nmost_common = action_counts.iloc[0]\nleast_common = action_counts.iloc[-1]\nimbalance_ratio = most_common / least_common if least_common > 0 else np.nan\n\nfig, ax = plt.subplots(1, 2, figsize=(14, 6), dpi=120)\n\n# Histogram of action frequencies (log scale)\nsns.histplot(action_counts, bins=40, log_scale=(True, False), color=PALETTE[\"deep_blue\"], ax=ax[0])\nax[0].set_title(\"Distribution of Action Frequencies\", color=PALETTE[\"slate_gray\"], fontsize=14, pad=12)\nax[0].set_xlabel(\"Instances per Action (log scale)\")\nax[0].set_ylabel(\"Number of Actions\")\nax[0].spines[\"top\"].set_visible(False)\nax[0].spines[\"right\"].set_visible(False)\nax[0].grid(axis=\"y\", linestyle=\"--\", alpha=0.35)\n\n# Bar plot for rare vs common classes\nax[1].bar([\"Rare (≤10)\", \"Common (≥1000)\"], [rare_count, common_count], color=[PALETTE[\"teal\"], PALETTE[\"deep_blue\"]])\nfor i, v in enumerate([rare_count, common_count]):\n ax[1].text(i, v + 1, f\"{v:,}\", ha=\"center\", va=\"bottom\", color=PALETTE[\"slate_gray\"], fontsize=11)\n\nax[1].set_title(\"Counts of Rare vs Common Actions\", color=PALETTE[\"slate_gray\"], fontsize=14, pad=12)\nax[1].set_ylabel(\"Number of Actions\")\nax[1].spines[\"top\"].set_visible(False)\nax[1].spines[\"right\"].set_visible(False)\nax[1].grid(axis=\"y\", linestyle=\"--\", alpha=0.35)\n\nplt.suptitle(\n f\"Class Imbalance Analysis\\nMost common: {most_common:,} • Least common: {least_common:,} • Ratio: {imbalance_ratio:.1f}x\",\n color=PALETTE[\"slate_gray\"], fontsize=13, y=1.02\n)\n\nplt.tight_layout()\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-04T13:21:36.401086Z","iopub.execute_input":"2025-10-04T13:21:36.401501Z","iopub.status.idle":"2025-10-04T13:21:37.584112Z","shell.execute_reply.started":"2025-10-04T13:21:36.401479Z","shell.execute_reply":"2025-10-04T13:21:37.583057Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if \"usable_train_rows\" not in globals():\n anno_root = folders[\"train_annotation\"]\n valid_annos = set()\n if anno_root.exists():\n for lab_dir in sorted(anno_root.iterdir()):\n if not lab_dir.is_dir():\n continue\n for f in lab_dir.glob(\"*.parquet\"):\n valid_annos.add((lab_dir.name, f.stem))\n usable_mask = train_csv.apply(lambda r: (r[\"lab_id\"], r[\"video_id\"]) in valid_annos, axis=1)\n usable_train_rows = train_csv[usable_mask]\n\nsample_labs = list(pd.unique(usable_train_rows[\"lab_id\"]))[:8]\nsample_videos = []\nfor lab in sample_labs:\n lab_videos = usable_train_rows[usable_train_rows[\"lab_id\"] == lab].head(3)\n sample_videos.extend(lab_videos[[\"lab_id\", \"video_id\"]].values.tolist())\n\ntracking_samples = []\nbodypart_inventory = {}\ncoordinate_stats = []\n\nfor lab_id, video_id in sample_videos:\n track_path = DATA / \"train_tracking\" / lab_id / f\"{video_id}.parquet\"\n if not track_path.exists():\n continue\n try:\n df_track = pd.read_parquet(track_path)\n df_track = df_track.head(1000).copy()\n df_track[\"lab_id\"] = lab_id\n df_track[\"video_id\"] = video_id\n tracking_samples.append(df_track)\n unique_bodyparts = df_track[\"bodypart\"].unique()\n bodypart_inventory[lab_id] = sorted(unique_bodyparts.tolist())\n coord_stats = df_track[[\"x\", \"y\"]].describe()\n coordinate_stats.append({\n \"lab_id\": lab_id,\n \"video_id\": video_id,\n \"x_min\": coord_stats.loc[\"min\", \"x\"],\n \"x_max\": coord_stats.loc[\"max\", \"x\"],\n \"y_min\": coord_stats.loc[\"min\", \"y\"],\n \"y_max\": coord_stats.loc[\"max\", \"y\"],\n \"x_mean\": coord_stats.loc[\"mean\", \"x\"],\n \"y_mean\": coord_stats.loc[\"mean\", \"y\"],\n \"sampled_frames\": len(df_track),\n \"unique_bodyparts\": len(unique_bodyparts)\n })\n except Exception:\n continue\n\nif len(tracking_samples) == 0:\n display(HTML(f\"
No tracking samples found for the selected labs/videos.
\"))\nelse:\n combined_tracking = pd.concat(tracking_samples, ignore_index=True)\n coord_df = pd.DataFrame(coordinate_stats).sort_values([\"lab_id\", \"video_id\"]).reset_index(drop=True)\n\n coord_df_display = coord_df.copy()\n coord_df_display[\"x_range\"] = (coord_df_display[\"x_max\"] - coord_df_display[\"x_min\"]).round(1)\n coord_df_display[\"y_range\"] = (coord_df_display[\"y_max\"] - coord_df_display[\"y_min\"]).round(1)\n coord_df_display = coord_df_display[[\n \"lab_id\", \"video_id\", \"sampled_frames\", \"unique_bodyparts\",\n \"x_min\", \"x_max\", \"x_range\", \"x_mean\",\n \"y_min\", \"y_max\", \"y_range\", \"y_mean\"\n ]]\n\n display(HTML(f\"\"\"\n
\n

Pose Data Sample — Coordinate Summary

\n
\n Total sampled frames: {len(combined_tracking):,}  • \n Labs sampled: {combined_tracking['lab_id'].nunique()}  • \n Videos sampled: {combined_tracking['video_id'].nunique()}  • \n Unique mice: {combined_tracking['mouse_id'].nunique()}\n
\n
\n \"\"\"))\n\n display(\n coord_df_display.style.format({\n \"sampled_frames\":\"{:,}\",\n \"unique_bodyparts\":\"{:,}\",\n \"x_min\":\"{:.1f}\", \"x_max\":\"{:.1f}\", \"x_range\":\"{:.1f}\", \"x_mean\":\"{:.1f}\",\n \"y_min\":\"{:.1f}\", \"y_max\":\"{:.1f}\", \"y_range\":\"{:.1f}\", \"y_mean\":\"{:.1f}\"\n }).set_table_styles([{\"selector\":\"th\",\"props\":[(\"background-color\",PALETTE[\"muted\"]),(\"color\",\"white\")]}])\n )","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-04T13:21:42.239948Z","iopub.execute_input":"2025-10-04T13:21:42.240271Z","iopub.status.idle":"2025-10-04T13:21:44.753717Z","shell.execute_reply.started":"2025-10-04T13:21:42.240248Z","shell.execute_reply":"2025-10-04T13:21:44.752653Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"all_bodyparts = set()\nfor bodyparts in bodypart_inventory.values():\n all_bodyparts.update(bodyparts)\n\nbodypart_matrix = pd.DataFrame(index=sorted(bodypart_inventory.keys()), columns=sorted(all_bodyparts))\nfor lab, bodyparts in bodypart_inventory.items():\n for bp in all_bodyparts:\n bodypart_matrix.loc[lab, bp] = 1 if bp in bodyparts else 0\nbodypart_matrix = bodypart_matrix.astype(int)\n\ntotal_bodyparts = len(all_bodyparts)\nper_lab_counts = bodypart_matrix.sum(axis=1)\nmin_bp, max_bp = int(per_lab_counts.min()), int(per_lab_counts.max())\n\nplt.figure(figsize=(14, max(4, 0.25 * bodypart_matrix.shape[0])), dpi=120)\ncmap = sns.color_palette([PALETTE[\"muted\"], PALETTE[\"teal\"]])\nax = sns.heatmap(\n bodypart_matrix,\n cmap=cmap,\n cbar=False,\n linewidths=0.25,\n linecolor=\"#E6EEF2\"\n)\nax.set_title(\n f\"Bodypart Tracking Overview — {total_bodyparts} unique bodyparts \"\n f\"(per lab range: {min_bp}–{max_bp})\",\n color=PALETTE[\"slate_gray\"], fontsize=14, pad=12\n)\nax.set_xlabel(\"Bodyparts\", color=PALETTE[\"slate_gray\"])\nax.set_ylabel(\"Labs\", color=PALETTE[\"slate_gray\"])\nax.tick_params(axis=\"x\", rotation=65)\nplt.tight_layout()\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-04T13:21:49.119529Z","iopub.execute_input":"2025-10-04T13:21:49.120009Z","iopub.status.idle":"2025-10-04T13:21:49.498745Z","shell.execute_reply.started":"2025-10-04T13:21:49.119984Z","shell.execute_reply":"2025-10-04T13:21:49.497551Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"common_bodyparts = bodypart_matrix.sum(axis=0).sort_values(ascending=False)\ntop15_bodyparts = common_bodyparts.head(15).reset_index()\ntop15_bodyparts.columns = [\"Bodypart\", \"Lab Count\"]\n\nplt.figure(figsize=(12, 6), dpi=120)\nax = sns.barplot(\n data=top15_bodyparts,\n x=\"Lab Count\", y=\"Bodypart\",\n palette=[PALETTE[\"teal\"]] * len(top15_bodyparts)\n)\n\nfor i, v in enumerate(top15_bodyparts[\"Lab Count\"]):\n ax.text(v + 0.3, i, f\"{v}\", va=\"center\", ha=\"left\", color=PALETTE[\"slate_gray\"], fontsize=10)\n\nax.set_title(\"Most Commonly Tracked Bodyparts (Top 15)\", color=PALETTE[\"slate_gray\"], fontsize=14, pad=12)\nax.set_xlabel(\"Number of Labs Tracking\")\nax.set_ylabel(\"\")\nax.spines[\"top\"].set_visible(False)\nax.spines[\"right\"].set_visible(False)\nax.grid(axis=\"x\", linestyle=\"--\", alpha=0.35)\nplt.tight_layout()\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-04T13:21:53.934036Z","iopub.execute_input":"2025-10-04T13:21:53.934352Z","iopub.status.idle":"2025-10-04T13:21:54.371895Z","shell.execute_reply.started":"2025-10-04T13:21:53.934330Z","shell.execute_reply":"2025-10-04T13:21:54.371004Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"plt.figure(figsize=(12, 7), dpi=120)\n\nfor i, row in coord_df.iterrows():\n plt.plot([row[\"x_min\"], row[\"x_max\"]], [i, i], color=PALETTE[\"deep_blue\"], linewidth=2)\n plt.scatter(row[\"x_min\"], i, color=PALETTE[\"teal\"], s=50, zorder=3)\n plt.scatter(row[\"x_max\"], i, color=PALETTE[\"teal\"], s=50, zorder=3)\n\nplt.yticks(range(len(coord_df)), coord_df[\"lab_id\"])\nplt.xlabel(\"X-coordinate Range (pixels)\", color=PALETTE[\"slate_gray\"])\nplt.title(\"Coordinate System Analysis — X-axis Ranges by Lab\", color=PALETTE[\"slate_gray\"], fontsize=14, pad=12)\nplt.grid(axis=\"x\", linestyle=\"--\", alpha=0.35)\nplt.tight_layout()\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-04T13:21:57.600937Z","iopub.execute_input":"2025-10-04T13:21:57.601253Z","iopub.status.idle":"2025-10-04T13:21:58.501235Z","shell.execute_reply.started":"2025-10-04T13:21:57.601228Z","shell.execute_reply":"2025-10-04T13:21:58.500260Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"fig, axes = plt.subplots(1, 3, figsize=(18, 8), dpi=120, constrained_layout=True)\n\n# 1) Arena centers by lab (mean x, mean y)\nax = axes[0]\npalette_lab = sns.color_palette(\"tab20\", n_colors=len(coord_df))\nsc = ax.scatter(coord_df[\"x_mean\"], coord_df[\"y_mean\"], c=range(len(coord_df)), cmap=\"tab20\", s=90, alpha=0.85)\nfor i, row in coord_df.reset_index().iterrows():\n ax.annotate(row[\"lab_id\"][:12], (row[\"x_mean\"], row[\"y_mean\"]), xytext=(6, 4), textcoords=\"offset points\",\n fontsize=8, color=PALETTE[\"slate_gray\"])\nax.set_xlabel(\"Mean X Coordinate\", color=PALETTE[\"slate_gray\"])\nax.set_ylabel(\"Mean Y Coordinate\", color=PALETTE[\"slate_gray\"])\nax.set_title(\"Arena Centers by Lab\", color=PALETTE[\"slate_gray\"], fontsize=13, pad=8)\n\n# 2) Arena dimensions by lab (range x, range y)\nax = axes[1]\nx_range = coord_df[\"x_max\"] - coord_df[\"x_min\"]\ny_range = coord_df[\"y_max\"] - coord_df[\"y_min\"]\nax.scatter(x_range, y_range, c=range(len(coord_df)), cmap=\"tab20\", s=90, alpha=0.85)\nfor i, row in coord_df.reset_index().iterrows():\n xr = row[\"x_max\"] - row[\"x_min\"]\n yr = row[\"y_max\"] - row[\"y_min\"]\n ax.annotate(row[\"lab_id\"][:10], (xr, yr), xytext=(6, 4), textcoords=\"offset points\", fontsize=8, color=PALETTE[\"slate_gray\"])\nax.set_xlabel(\"X Range (px)\", color=PALETTE[\"slate_gray\"])\nax.set_ylabel(\"Y Range (px)\", color=PALETTE[\"slate_gray\"])\nax.set_title(\"Arena Dimensions by Lab\", color=PALETTE[\"slate_gray\"], fontsize=13, pad=8)\n\n# 3) Coordinate distribution sample (pose points colored by lab)\nax = axes[2]\nsample_n = min(12000, len(combined_tracking))\nsample_df = combined_tracking.sample(sample_n, random_state=42)\nlabs = sample_df[\"lab_id\"].unique()\npal = sns.color_palette(\"tab20\", n_colors=len(labs))\nlab_to_color = {lab: pal[i % len(pal)] for i, lab in enumerate(labs)}\nax.scatter(sample_df[\"x\"], sample_df[\"y\"], s=6,\n c=[lab_to_color[l] for l in sample_df[\"lab_id\"].values], alpha=0.6)\nax.invert_yaxis()\nax.set_xlabel(\"x (pixels)\", color=PALETTE[\"slate_gray\"])\nax.set_ylabel(\"y (pixels)\", color=PALETTE[\"slate_gray\"])\nax.set_title(\"Coordinate Distribution Sample (colored by lab)\", color=PALETTE[\"slate_gray\"], fontsize=13, pad=8)\n# compact legend: show only top 8 labs to avoid clutter\ntop_labs_for_legend = list(pd.Series(sample_df[\"lab_id\"]).value_counts().head(8).index)\nhandles = [plt.Line2D([], [], marker=\"o\", linestyle=\"\", color=lab_to_color[l], markersize=6) for l in top_labs_for_legend]\nax.legend(handles, top_labs_for_legend, title=\"Top labs\", bbox_to_anchor=(1.02, 1), loc=\"upper left\", frameon=False)\n\nplt.suptitle(\"Coordinate System Summary — centers, dimensions, and sampled point cloud\", color=PALETTE[\"slate_gray\"], fontsize=14, y=1.02)\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-04T13:22:02.738418Z","iopub.execute_input":"2025-10-04T13:22:02.738770Z","iopub.status.idle":"2025-10-04T13:22:04.286740Z","shell.execute_reply.started":"2025-10-04T13:22:02.738742Z","shell.execute_reply":"2025-10-04T13:22:04.285721Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if \"sample_videos\" not in globals():\n sample_labs = list(pd.unique(usable_train_rows[\"lab_id\"]))[:8]\n sample_videos = []\n for lab in sample_labs:\n lab_videos = usable_train_rows[usable_train_rows[\"lab_id\"] == lab].head(3)\n sample_videos.extend(lab_videos[[\"lab_id\", \"video_id\"]].values.tolist())\n\ntemporal_analysis = []\nfor lab_id, video_id in sample_videos[:10]:\n track_path = DATA / \"train_tracking\" / lab_id / f\"{video_id}.parquet\"\n if not track_path.exists():\n continue\n try:\n df = pd.read_parquet(track_path)\n df_pivot = df.pivot_table(index=\"video_frame\", columns=[\"mouse_id\", \"bodypart\"], values=[\"x\", \"y\"])\n if len(df_pivot) > 100:\n sample_coords = df_pivot.iloc[:100, :4].values.astype(float)\n velocities = np.diff(sample_coords, axis=0)\n vel_magnitude = np.sqrt(np.nansum(velocities**2, axis=1))\n temporal_analysis.append({\n \"lab_id\": lab_id,\n \"video_id\": video_id,\n \"mean_velocity\": float(np.nanmean(vel_magnitude)),\n \"max_velocity\": float(np.nanmax(vel_magnitude)),\n \"velocity_std\": float(np.nanstd(vel_magnitude))\n })\n except Exception:\n continue\n\ntemporal_df = pd.DataFrame(temporal_analysis)\nif temporal_df.empty:\n raise ValueError(\"No temporal samples found for the chosen videos. Re-run sampling or choose different videos.\")\n\ntemporal_df = temporal_df.round({\n \"mean_velocity\": 2,\n \"max_velocity\": 2,\n \"velocity_std\": 2\n}).sort_values([\"mean_velocity\", \"max_velocity\"], ascending=[False, False]).reset_index(drop=True)\n\ndisplay(\n temporal_df.style.background_gradient(\n subset=[\"mean_velocity\", \"max_velocity\", \"velocity_std\"], cmap=\"YlGnBu\"\n ).set_caption(\"Temporal Consistency Analysis — Velocity Statistics\")\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-04T13:25:38.346495Z","iopub.execute_input":"2025-10-04T13:25:38.346896Z","iopub.status.idle":"2025-10-04T13:25:50.707571Z","shell.execute_reply.started":"2025-10-04T13:25:38.346873Z","shell.execute_reply":"2025-10-04T13:25:50.706749Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if \"missing_summary\" not in globals():\n # ensure usable_train_rows exists\n if \"usable_train_rows\" not in globals():\n anno_root = folders[\"train_annotation\"]\n valid_annos = set()\n if anno_root.exists():\n for lab_dir in sorted(anno_root.iterdir()):\n if not lab_dir.is_dir():\n continue\n for f in lab_dir.glob(\"*.parquet\"):\n valid_annos.add((lab_dir.name, f.stem))\n usable_mask = train_csv.apply(lambda r: (r[\"lab_id\"], r[\"video_id\"]) in valid_annos, axis=1)\n usable_train_rows = train_csv[usable_mask]\n\n # ensure combined_tracking exists (sample if necessary)\n if \"combined_tracking\" not in globals():\n sample_labs = list(pd.unique(usable_train_rows[\"lab_id\"]))[:8]\n sample_videos = []\n for lab in sample_labs:\n lab_videos = usable_train_rows[usable_train_rows[\"lab_id\"] == lab].head(3)\n sample_videos.extend(lab_videos[[\"lab_id\", \"video_id\"]].values.tolist())\n\n tracking_samples = []\n for lab_id, video_id in sample_videos:\n track_path = DATA / \"train_tracking\" / lab_id / f\"{video_id}.parquet\"\n if not track_path.exists():\n continue\n try:\n df_track = pd.read_parquet(track_path)\n df_track = df_track.head(1000).copy()\n df_track[\"lab_id\"] = lab_id\n df_track[\"video_id\"] = video_id\n tracking_samples.append(df_track)\n except Exception:\n continue\n\n if len(tracking_samples) == 0:\n raise RuntimeError(\"No tracking samples found to compute missing_summary. Ensure tracking files are available.\")\n combined_tracking = pd.concat(tracking_samples, ignore_index=True)\n\n # groupby + apply with include_groups=False to silence deprecation warning\n missing_analysis = combined_tracking.groupby(['lab_id', 'mouse_id', 'bodypart']).apply(\n lambda x: pd.Series({\n 'total_frames': len(x),\n 'missing_x': int(x['x'].isnull().sum()),\n 'missing_y': int(x['y'].isnull().sum()),\n 'missing_either': int((x['x'].isnull() | x['y'].isnull()).sum())\n }),\n include_groups=False\n ).reset_index()\n\n # aggregate to lab-level\n missing_summary = missing_analysis.groupby('lab_id').agg({\n 'missing_either': ['sum', 'mean'],\n 'total_frames': 'sum'\n }).round(3)\n missing_summary.columns = ['total_missing', 'avg_missing_rate', 'total_frames']\n missing_summary['missing_percentage'] = (missing_summary['total_missing'] / missing_summary['total_frames'] * 100).round(2)\n missing_summary = missing_summary.reset_index()\n\n# compute KPIs using coord_df and missing_summary (assumes coord_df exists)\ncoord_range_variation = coord_df[\"x_max\"].max() - coord_df[\"x_min\"].min()\narena_cv = ((coord_df[\"x_max\"] - coord_df[\"x_min\"]).std() /\n (coord_df[\"x_max\"] - coord_df[\"x_min\"]).mean()) if (coord_df[\"x_max\"] - coord_df[\"x_min\"]).mean() != 0 else np.nan\nmissing_min, missing_max = missing_summary[\"missing_percentage\"].min(), missing_summary[\"missing_percentage\"].max()\n\nfrom IPython.display import HTML, display\ndisplay(HTML(f\"\"\"\n
\n

Summary Statistics

\n \n
\n\"\"\"))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-04T13:25:57.477534Z","iopub.execute_input":"2025-10-04T13:25:57.478216Z","iopub.status.idle":"2025-10-04T13:25:57.609291Z","shell.execute_reply.started":"2025-10-04T13:25:57.478187Z","shell.execute_reply":"2025-10-04T13:25:57.608469Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Video Simulation","metadata":{}},{"cell_type":"code","source":"def create_solution_df(dataset):\n \"\"\"Create the solution dataframe for validating out-of-fold predictions.\n\n From https://www.kaggle.com/code/ambrosm/mabe-validated-baseline-without-machine-learning/\n \n Parameters:\n dataset: (a subset of) the train dataframe\n \n Return values:\n solution: solution dataframe in the correct format for the score() function\n \"\"\"\n solution = []\n for train_idx, row in tqdm(dataset.iterrows(), total=len(dataset)):\n \n # Load annotation file\n lab_id = row['lab_id']\n if lab_id.startswith('MABe22'): continue\n video_id = row['video_id']\n path = f\"/kaggle/input/MABe-mouse-behavior-detection/train_annotation/{lab_id}/{video_id}.parquet\"\n try:\n annot = pd.read_parquet(path.replace('train_tracking', 'train_annotation'))\n except FileNotFoundError:\n print(f\"No annotations for {path}\")\n continue\n \n # Add all annotations to the solution\n self_annotation = annot.target_id == annot.agent_id\n annot['lab_id'] = lab_id\n annot['video_id'] = video_id\n annot['behaviors_labeled'] = row['behaviors_labeled']\n annot['agent_id'] = annot['agent_id'].apply(lambda s: f\"mouse{s}\")\n annot['target_id'] = np.where(self_annotation, 'self', annot['target_id'].apply(lambda s: f\"mouse{s}\"))\n solution.append(annot)\n \n solution = pd.concat(solution)\n return solution\n\nsolution = create_solution_df(train_subset)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-04T13:28:50.565629Z","iopub.execute_input":"2025-10-04T13:28:50.566009Z","iopub.status.idle":"2025-10-04T13:28:55.262517Z","shell.execute_reply.started":"2025-10-04T13:28:50.565982Z","shell.execute_reply":"2025-10-04T13:28:55.261542Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\nimport os\nfrom pathlib import Path\nfrom tqdm import tqdm\nfrom IPython.display import Video, display\n\ndataset_path = Path('/kaggle/input/MABe-mouse-behavior-detection')\ntrain = pd.read_csv(dataset_path / 'train.csv')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-04T13:29:01.945251Z","iopub.execute_input":"2025-10-04T13:29:01.945589Z","iopub.status.idle":"2025-10-04T13:29:02.022370Z","shell.execute_reply.started":"2025-10-04T13:29:01.945562Z","shell.execute_reply":"2025-10-04T13:29:02.021498Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os, json, math, tempfile, subprocess, base64\nfrom functools import lru_cache\nfrom typing import Optional, List, Tuple, Dict\nimport numpy as np, pandas as pd, cv2\nfrom IPython.display import Video\nfrom tqdm.auto import tqdm\n\ndef distinct_colors_bgr(n, sat=200, val=235, hue_offset=0):\n if n <= 0: return []\n hues = (np.linspace(0,179,n,endpoint=False)+hue_offset)%180\n hsv = np.stack([hues, np.full(n,sat), np.full(n,val)],1).astype(np.uint8)[None,...]\n return [tuple(map(int,c)) for c in cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)[0]]\n\ndef _extract_bps(meta: Optional[pd.DataFrame]):\n if meta is None or meta.empty: return None\n row = meta.iloc[0]\n for col in [\"body_parts_tracked\",\"bodyparts_tracked\",\"body_parts\",\"tracked_body_parts\"]:\n if col in meta.columns:\n v = row[col]\n if isinstance(v,list): return v\n if isinstance(v,str):\n try:\n p = json.loads(v)\n if isinstance(p,list): return p\n except: pass\n return None\n\ndef _edges(available: List[str]):\n cand=[(\"nose\",\"neck\"),(\"neck\",\"ear_left\"),(\"neck\",\"ear_right\"),\n (\"neck\",\"hip_left\"),(\"neck\",\"hip_right\"),\n (\"hip_left\",\"tail_base\"),(\"hip_right\",\"tail_base\"),\n (\"body_center\",\"neck\"),(\"body_center\",\"nose\"),(\"body_center\",\"tail_base\")]\n seen=set(); out=[]\n for a,b in cand:\n if a in available and b in available:\n k=tuple(sorted((a,b)))\n if k not in seen: seen.add(k); out.append((a,b))\n return out\n\ndef export_square_nocrop_mp4_embed(\n df_episode: pd.DataFrame,\n df_episode_meta: Optional[pd.DataFrame]=None,\n df_annotation: Optional[pd.DataFrame]=None, # <— NEW\n *,\n size_px:int=640, out_fps:Optional[int]=None,\n frame_start:Optional[int]=None, frame_end:Optional[int]=None, frame_stride:int=1,\n show_skeleton:bool=True, show_trails:bool=False, trail_len:int=25,\n bodyparts:Optional[List[str]]=None, marker_radius:int=3, border_thickness:int=2, line_thickness:int=2,\n # action caption options (NEW)\n show_actions: bool=True,\n action_anchor_priority: Tuple[str,...]=(\"nose\",\"neck\",\"body_center\"), # where to place label\n action_font_scale: float=0.5,\n action_thickness: int=1,\n action_text_color: Tuple[int,int,int]=(0,0,0), # BGR\n action_box_bg: Tuple[int,int,int]=(255,255,255), # BGR\n action_box_pad: int=3,\n action_box_border: int=1,\n action_box_max_width_ratio: float=0.9, # clamp box within canvas\n action_joiner: str=\" | \", # if multiple concurrent actions\n # output / encoding\n out_path:str=\"episode_square.mp4\",\n ffmpeg_preset:str=\"veryfast\", ffmpeg_crf:int=23,\n return_embed:bool=True, max_embed_mb:int=120, display_width:int=720\n):\n # basic validation\n need={\"video_frame\",\"mouse_id\",\"bodypart\",\"x\",\"y\"}\n miss=need-set(df_episode.columns)\n if miss: raise ValueError(f\"df_episode missing: {sorted(miss)}\")\n df=df_episode.sort_values([\"video_frame\",\"mouse_id\",\"bodypart\"]).reset_index(drop=True)\n\n # source plane (no crop)\n if df_episode_meta is not None and not df_episode_meta.empty:\n r=df_episode_meta.iloc[0]\n W=r.get(\"video_width_pix\",np.nan); H=r.get(\"video_height_pix\",np.nan)\n if np.isfinite(W) and np.isfinite(H):\n xmin=ymin=0.0; src_w=float(W); src_h=float(H)\n else:\n xmin, xmax = float(df.x.min()), float(df.x.max())\n ymin, ymax = float(df.y.min()), float(df.y.max())\n src_w=max(1.0,xmax-xmin); src_h=max(1.0,ymax-ymin)\n else:\n xmin, xmax = float(df.x.min()), float(df.x.max())\n ymin, ymax = float(df.y.min()), float(df.y.max())\n src_w=max(1.0,xmax-xmin); src_h=max(1.0,ymax-ymin); xmin=ymin=0.0\n\n N=int(size_px); s=min(N/src_w, N/src_h); pad_x=0.5*(N-s*src_w); pad_y=0.5*(N-s*src_h)\n def map_xy(x,y):\n X=pad_x + s*(float(x)-xmin); Y=pad_y + s*(float(y)-ymin)\n return max(0,min(N-1,int(round(X)))), max(0,min(N-1,int(round(Y))))\n\n # fps\n if out_fps is None:\n v = df_episode_meta.iloc[0][\"frames_per_second\"] if (df_episode_meta is not None and not df_episode_meta.empty and \"frames_per_second\" in df_episode_meta.columns) else None\n out_fps = int(v) if (v is not None and pd.notna(v)) else 12\n\n # frames\n all_frames=np.sort(df.video_frame.dropna().unique())\n if all_frames.size==0: raise ValueError(\"No frames.\")\n if frame_start is None: frame_start=int(all_frames[0])\n if frame_end is None: frame_end =int(all_frames[-1])\n frames=[int(f) for f in all_frames if frame_start<=f<=frame_end][::max(1,int(frame_stride))]\n if not frames: raise ValueError(\"No frames after filtering.\")\n frames_np = np.array(frames, dtype=int)\n\n # bodyparts/colors\n data_bps=sorted(df.bodypart.dropna().unique().tolist())\n meta_bps=_extract_bps(df_episode_meta)\n available_bps=meta_bps if meta_bps else data_bps\n if bodyparts is None: bodyparts=available_bps\n else: bodyparts=[bp for bp in bodyparts if bp in available_bps]\n edges=_edges(bodyparts) if show_skeleton else []\n bp_cols=distinct_colors_bgr(len(available_bps), sat=200, val=235, hue_offset=0)\n bp_to_bgr={bp: bp_cols[i] for i,bp in enumerate(available_bps)}\n mice=sorted(df.mouse_id.dropna().unique().tolist())\n mouse_cols=distinct_colors_bgr(len(mice), sat=240, val=220, hue_offset=17)\n mouse_to_bgr={mid: mouse_cols[i] for i,mid in enumerate(mice)}\n\n # --- Build fast lookup for annotations: (frame_idx, agent_id) -> [actions]\n ann_map: Dict[Tuple[int,int], List[str]] = {}\n \n if show_actions and df_annotation is not None and not df_annotation.empty:\n req_cols = {\"agent_id\",\"action\",\"start_frame\",\"stop_frame\"}\n miss = req_cols - set(df_annotation.columns)\n if miss:\n raise ValueError(f\"df_annotation missing: {sorted(miss)}\")\n # keep only agents in this episode\n ann = df_annotation.copy()\n ann = ann[ann[\"agent_id\"].isin(mice)]\n # normalize numeric\n for c in (\"start_frame\",\"stop_frame\"):\n ann[c] = pd.to_numeric(ann[c], errors=\"coerce\").astype(\"Int64\")\n ann = ann.dropna(subset=[\"start_frame\",\"stop_frame\",\"action\",\"agent_id\"])\n if not ann.empty:\n ann[\"start_frame\"] = ann[\"start_frame\"].astype(int)\n ann[\"stop_frame\"] = ann[\"stop_frame\"].astype(int)\n if ann[\"start_frame\"].gt(ann[\"stop_frame\"]).any():\n # swap where needed\n bad = ann[\"start_frame\"] > ann[\"stop_frame\"]\n tmp = ann.loc[bad, \"start_frame\"].values\n ann.loc[bad, \"start_frame\"] = ann.loc[bad, \"stop_frame\"].values\n ann.loc[bad, \"stop_frame\"] = tmp\n # For each annotation row, add entries for intersecting frames (use binary search into frames)\n for row in ann.itertuples(index=False):\n a0, a1 = int(row.start_frame), int(row.stop_frame)\n # overlap with rendered frames\n i0 = int(np.searchsorted(frames_np, a0, side=\"left\"))\n i1 = int(np.searchsorted(frames_np, a1, side=\"right\"))\n if i0 >= i1: \n continue\n agent = int(row.agent_id) if pd.api.types.is_numeric_dtype(type(row.agent_id)) or isinstance(row.agent_id, (int,np.integer)) else row.agent_id\n act = str(row.action)\n for f in frames_np[i0:i1]:\n ann_map.setdefault((int(f), agent), []).append(act)\n\n @lru_cache(None)\n def fm(frame, mid):\n sub=df[(df.video_frame==frame)&(df.mouse_id==mid)]\n return sub[sub.bodypart.isin(bodyparts)]\n\n trails={mid:[] for mid in mice}\n\n # write tmp with OpenCV, then ffmpeg → H.264\n tmp_output_path = \"tmp_\" + out_path\n vw=cv2.VideoWriter(tmp_output_path, cv2.VideoWriter_fourcc(*\"mp4v\"), out_fps, (N,N))\n if not vw.isOpened(): raise RuntimeError(\"OpenCV VideoWriter failed.\")\n\n try:\n for f in tqdm(frames, total=len(frames), desc=\"Rendering frames\", unit=\"f\",\n smoothing=0.1, mininterval=0.1, leave=False):\n canvas=np.full((N,N,3),(255,255,255),dtype=np.uint8) # white background\n for mid in mice:\n sub=fm(f,mid)\n if sub.empty: \n # even if no body part this frame, we might still want to write the action near last trail point (optional).\n continue\n bp_map={}\n for r in sub.itertuples(index=False):\n x,y=map_xy(r.x,r.y); bp_map[r.bodypart]=(x,y)\n # skeleton\n if show_skeleton:\n for p1,p2 in edges:\n if p1 in bp_map and p2 in bp_map:\n x1,y1=bp_map[p1]; x2,y2=bp_map[p2]\n cv2.line(canvas,(x1,y1),(x2,y2), bp_to_bgr.get(p1,(0,0,0)), thickness=line_thickness, lineType=cv2.LINE_AA)\n # joints\n for r in sub.itertuples(index=False):\n x,y=bp_map[r.bodypart]\n cv2.circle(canvas,(x,y),marker_radius, bp_to_bgr.get(r.bodypart,(0,0,0)), -1, cv2.LINE_AA)\n cv2.circle(canvas,(x,y),marker_radius, mouse_to_bgr.get(mid,(0,0,0)), border_thickness, cv2.LINE_AA)\n\n # trails\n ref = None\n for k in action_anchor_priority:\n if k in bp_map: \n ref = bp_map[k]; break\n if ref is None and bp_map:\n xs,ys=zip(*bp_map.values()); ref=(int(round(np.mean(xs))), int(round(np.mean(ys))))\n\n if show_trails and ref is not None:\n t=trails[mid]; t.append(ref); \n if len(t)>trail_len: trails[mid]=t[-trail_len:]\n pts=np.array(trails[mid],np.int32).reshape(-1,1,2)\n cv2.polylines(canvas,[pts],False, mouse_to_bgr[mid], 1, cv2.LINE_AA)\n\n # action captions\n if show_actions and ref is not None:\n acts = ann_map.get((int(f), mid))\n if acts:\n # make unique but stable\n uniq = list(dict.fromkeys(a.strip() for a in acts if str(a).strip()))\n if uniq:\n label = action_joiner.join(uniq)\n # measure text\n font = cv2.FONT_HERSHEY_SIMPLEX\n (tw, th), baseline = cv2.getTextSize(label, font, action_font_scale, action_thickness)\n pad = int(action_box_pad)\n box_w = min(tw + 2*pad, int(action_box_max_width_ratio * N))\n # position the box above the anchor, clamp within canvas\n x0 = int(ref[0] - box_w//2)\n y0 = int(ref[1] - marker_radius - 6 - th - baseline - 2*pad)\n x0 = max(0, min(N - box_w, x0))\n y0 = max(0, y0)\n # background box\n x1 = x0 + box_w\n y1 = y0 + th + baseline + 2*pad\n cv2.rectangle(canvas, (x0,y0), (x1,y1), action_box_bg, thickness=-1)\n if action_box_border > 0:\n cv2.rectangle(canvas, (x0,y0), (x1,y1), mouse_to_bgr.get(mid,(0,0,0)), thickness=action_box_border, lineType=cv2.LINE_AA)\n # text (left padded)\n tx = x0 + pad\n ty = y1 - baseline - pad\n cv2.putText(canvas, label, (tx,ty), font, action_font_scale, action_text_color, action_thickness, cv2.LINE_AA)\n\n # Draw frame number (top-left corner)\n font = cv2.FONT_HERSHEY_SIMPLEX\n text = f\"Frame {f}\"\n scale = 0.5\n thickness = 1\n color = (0, 0, 0) # black\n margin = 50\n \n (text_w, text_h), baseline = cv2.getTextSize(text, font, scale, thickness)\n x = margin\n y = margin\n \n cv2.putText(canvas, text, (x, y), font, scale, color, thickness, cv2.LINE_AA)\n\n vw.write(canvas)\n finally:\n vw.release()\n\n # Re-encode with libx264 baseline/yuv420p-ish quality settings\n subprocess.run(\n [\"ffmpeg\", \"-y\", \"-i\", tmp_output_path, \"-crf\", str(ffmpeg_crf), \"-preset\", ffmpeg_preset, \"-vcodec\", \"libx264\", \"-pix_fmt\", \"yuv420p\", out_path],\n check=False,\n stdout=subprocess.DEVNULL,\n stderr=subprocess.DEVNULL\n )\n try:\n os.remove(tmp_output_path)\n except OSError:\n pass\n\n return out_path","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-04T13:29:04.447718Z","iopub.execute_input":"2025-10-04T13:29:04.448038Z","iopub.status.idle":"2025-10-04T13:29:04.496744Z","shell.execute_reply.started":"2025-10-04T13:29:04.447998Z","shell.execute_reply":"2025-10-04T13:29:04.495711Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from IPython.display import Video, display\n\ndef create_video(video_idx):\n if video_idx is None or not (0 <= int(video_idx) < len(train)):\n raise IndexError(f\"video_idx must be in [0, {len(train)-1}]\")\n \n # Example: iterate over *all* episodes described in the meta dataframe\n row = train.iloc[int(video_idx)]\n \n lab = str(row[\"lab_id\"])\n vid = str(int(row[\"video_id\"])) # cast to int then str for clean filename\n\n # Build the parquet file path\n pathto_tracking_data = os.path.join(\n dataset_path,\n \"train_tracking\",\n lab,\n f\"{vid}.parquet\"\n )\n\n print(f\"Episode Tracking {video_idx} → {pathto_tracking_data}\")\n\n # Load the tracking dataframe\n df_tracking = pd.read_parquet(pathto_tracking_data)\n\n # Build the parquet file path\n pathto_annotation_data = os.path.join(\n dataset_path,\n \"train_annotation\",\n lab,\n f\"{vid}.parquet\"\n )\n\n print(f\"Episode Annotation {video_idx} → {pathto_annotation_data}\")\n\n # Load the tracking dataframe\n df_annotation = pd.read_parquet(pathto_annotation_data) \n\n vid = export_square_nocrop_mp4_embed(\n df_tracking, train.iloc[[video_idx]], df_annotation,\n size_px=640, frame_stride=2, show_trails=False,\n out_path=f\"episode-{video_idx}.mp4\", # you also keep a file\n return_embed=True, max_embed_mb=120, display_width=720\n )\n\n display(Video(data=vid,\n embed=True,\n height=int(640),\n width=int(640))\n )\n\n return df_annotation # For debugging purposes","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-04T13:29:11.610891Z","iopub.execute_input":"2025-10-04T13:29:11.611584Z","iopub.status.idle":"2025-10-04T13:29:11.619174Z","shell.execute_reply.started":"2025-10-04T13:29:11.611553Z","shell.execute_reply":"2025-10-04T13:29:11.618211Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Set an index and run the cell(or all if not run before)","metadata":{}},{"cell_type":"code","source":"index = 0\nannot = create_video(index)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-04T13:29:15.324466Z","iopub.execute_input":"2025-10-04T13:29:15.324825Z","iopub.status.idle":"2025-10-04T13:33:09.623939Z","shell.execute_reply.started":"2025-10-04T13:29:15.324799Z","shell.execute_reply":"2025-10-04T13:33:09.622617Z"}},"outputs":[],"execution_count":null}]}