{"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":13719397,"sourceType":"competition"}],"dockerImageVersionId":31089,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# MABe Challenge - Social Action Recognition in Mice: Nearest neighbors\n\nThis is the original notebook for social action recognition with nearest neighbors. I've tried to explain what the code does—feel free to ask questions.\n\nThe notebook shows how to overcome the five challenges of this competition:\n1. Modeling for variable-size sets of mice\n2. Multiclass prediction with missing labels\n3. Transforming coordinates to an invariant representation\n4. A dataset that doesn't fit into memory\n5. Modeling for variable sets of body parts\n\nThe title of the notebook mentions *Nearest Neighbors* because in earlier versions I used nearest neighbors classification, an algorithm which doesn't need a lot of tuning. The current version uses LightGBM, and maybe I'll ensemble the two later.\n\nReferences\n- Competition: [MABe Challenge - Social Action Recognition in Mice](https://www.kaggle.com/competitions/MABe-mouse-behavior-detection)\n- [MABe EDA which makes sense ⭐️⭐️⭐️⭐️⭐️](https://www.kaggle.com/code/ambrosm/mabe-eda-which-makes-sense)\n- [MABe Validated baseline without machine learning](https://www.kaggle.com/code/ambrosm/mabe-validated-baseline-without-machine-learning)\n\nThis notebook can be run in validate or submission mode. If you look at other saved versions of this notebook, you'll see both modes. You can switch between the modes by setting the variable `validate_or_submit`:","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19"}},{"cell_type":"code","source":"validate_or_submit = 'validate' # 'validate' or 'submit' or 'stresstest'\nverbose = True","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-27T10:57:00.872255Z","iopub.execute_input":"2025-09-27T10:57:00.872544Z","iopub.status.idle":"2025-09-27T10:57:00.880553Z","shell.execute_reply.started":"2025-09-27T10:57:00.872521Z","shell.execute_reply":"2025-09-27T10:57:00.879551Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nimport matplotlib.pyplot as plt\nfrom tqdm import trange, tqdm\nimport itertools\nimport warnings\nimport json\nimport os\nimport lightgbm\n\nfrom sklearn.base import ClassifierMixin, BaseEstimator, clone\nfrom sklearn.model_selection import cross_val_predict, GroupKFold, train_test_split\nfrom sklearn.pipeline import make_pipeline\nfrom sklearn.impute import SimpleImputer\nfrom sklearn.preprocessing import StandardScaler\nfrom sklearn.neighbors import KNeighborsClassifier\nfrom sklearn.metrics import f1_score\n","metadata":{"trusted":true,"_kg_hide-input":true,"execution":{"iopub.status.busy":"2025-09-27T10:57:00.882120Z","iopub.execute_input":"2025-09-27T10:57:00.882842Z","iopub.status.idle":"2025-09-27T10:57:04.998004Z","shell.execute_reply.started":"2025-09-27T10:57:00.882810Z","shell.execute_reply":"2025-09-27T10:57:04.997131Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class TrainOnSubsetClassifier(ClassifierMixin, BaseEstimator):\n    \"\"\"Fit estimator to a subset of the training data.\"\"\"\n    def __init__(self, estimator, n_samples):\n        self.estimator = estimator\n        self.n_samples = n_samples\n\n    def fit(self, X, y):\n        downsample = len(X) // self.n_samples\n        downsample = max(downsample, 1)\n        self.estimator.fit(np.array(X, copy=False)[::downsample],\n                           np.array(y, copy=False)[::downsample])\n        self.classes_ = self.estimator.classes_\n        return self\n\n    def predict_proba(self, X):\n        if len(self.classes_) == 1:\n            return np.full((len(X), 1), 1.0)\n        probs = self.estimator.predict_proba(np.array(X))\n        return probs\n        \n    def predict(self, X):\n        return self.estimator.predict(np.array(X))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-27T10:57:04.998975Z","iopub.execute_input":"2025-09-27T10:57:04.999611Z","iopub.status.idle":"2025-09-27T10:57:05.007230Z","shell.execute_reply.started":"2025-09-27T10:57:04.999587Z","shell.execute_reply":"2025-09-27T10:57:05.006330Z"},"_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                # print(f'ignoring {video}', ','.join([str(row['agent_id']), str(row['target_id']), row['action']]), 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        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-09-27T10:57:05.009698Z","iopub.execute_input":"2025-09-27T10:57:05.009992Z","iopub.status.idle":"2025-09-27T10:57:05.335240Z","shell.execute_reply.started":"2025-09-27T10:57:05.009940Z","shell.execute_reply":"2025-09-27T10:57:05.334324Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"We start by reading the training metadata from train.csv.","metadata":{}},{"cell_type":"code","source":"train = pd.read_csv('/kaggle/input/MABe-mouse-behavior-detection/train.csv')\ntrain['n_mice'] = 4 - train[['mouse1_strain', 'mouse2_strain', 'mouse3_strain', 'mouse4_strain']].isna().sum(axis=1)\ntrain_without_mabe22 = train.query(\"~ lab_id.str.startswith('MABe22_')\")\n\ntest = pd.read_csv('/kaggle/input/MABe-mouse-behavior-detection/test.csv')\n\n# labs = list(np.unique(train.lab_id))\n\nbody_parts_tracked_list = list(np.unique(train.body_parts_tracked))\n\n# behaviors = list(train.behaviors_labeled.drop_duplicates().dropna())\n# behaviors = sorted(list({b.replace(\"'\", \"\") for bb in behaviors for b in json.loads(bb)}))\n# behaviors = [b.split(',') for b in behaviors]\n# behaviors = pd.DataFrame(behaviors, columns=['agent', 'target', 'action'])\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-27T10:57:05.336078Z","iopub.execute_input":"2025-09-27T10:57:05.336315Z","iopub.status.idle":"2025-09-27T10:57:05.527731Z","shell.execute_reply.started":"2025-09-27T10:57:05.336296Z","shell.execute_reply":"2025-09-27T10:57:05.526903Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"We collect all annotation files and bring the true labels into the format required by the competition scoring function.","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 _, 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)\n        except FileNotFoundError:\n            # MABe22 and one more training file lack annotations.\n            if verbose: print(f\"No annotations for {path}\")\n            continue\n    \n        # Add all annotations to the solution\n        annot['lab_id'] = lab_id\n        annot['video_id'] = video_id\n        annot['behaviors_labeled'] = row['behaviors_labeled']\n        annot['target_id'] = np.where(annot.target_id != annot.agent_id, annot['target_id'].apply(lambda s: f\"mouse{s}\"), 'self')\n        annot['agent_id'] = annot['agent_id'].apply(lambda s: f\"mouse{s}\")\n        solution.append(annot)\n    \n    solution = pd.concat(solution)\n    return solution\n\nif validate_or_submit == 'validate':\n    solution = create_solution_df(train_without_mabe22)","metadata":{"trusted":true,"_kg_hide-input":true,"execution":{"iopub.status.busy":"2025-09-27T10:57:05.528725Z","iopub.execute_input":"2025-09-27T10:57:05.529013Z","iopub.status.idle":"2025-09-27T10:57:12.934991Z","shell.execute_reply.started":"2025-09-27T10:57:05.528981Z","shell.execute_reply":"2025-09-27T10:57:12.933706Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Stress testing with unusual inputs\n\nAfter submission, this notebook will see a test set that it has never seen before. If the notebook crashes, debugging will be hard. It's better to stress-test the notebook before the submission by giving it some unusual inputs. The following hidden cell generate synthetic data with missing values, excessively long videos and so on.","metadata":{}},{"cell_type":"code","source":"if validate_or_submit == 'stresstest':\n    n_videos_per_lab = 2\n    \n    try:\n        os.mkdir(f\"stresstest_tracking\")\n    except FileExistsError:\n        pass\n    \n    rng = np.random.default_rng()\n    stresstest = pd.concat(\n        [train.query(\"video_id == 1459695188\")] # long video from BoisterousParrot\n        + [df.sample(min(n_videos_per_lab, len(df)), random_state=1) for (_, df) in train.groupby('lab_id')])\n    for _, row in tqdm(stresstest.iterrows(), total=len(stresstest)):\n        lab_id = row['lab_id']\n        video_id = row['video_id']\n        \n        # Load video\n        path = f\"/kaggle/input/MABe-mouse-behavior-detection/train_tracking/{lab_id}/{video_id}.parquet\"\n        vid = pd.read_parquet(path)\n    \n        if video_id == 1459695188: # long video from BoisterousParrot\n            vid = pd.concat([vid] * 3) # provoke out of memory (5 is too much)\n            vid['video_frame'] = np.arange(len(vid))\n    \n        # Drop some complete frames\n        dropped_frames = list(rng.choice(np.unique(vid.video_frame), size=100, replace=False))\n        vid = vid.query(\"~ video_frame.isin(@dropped_frames)\")\n        \n        # Drop a complete bodypart\n        if rng.uniform() < 0.2:\n            dropped_bodypart = rng.choice(np.unique(vid.bodypart), size=1, replace=False)[0]\n            vid = vid.query(\"bodypart != @dropped_bodypart\")\n        \n        # Drop a mouse\n        if rng.uniform() < 0.1:\n            vid = vid.query(\"mouse_id != 1\")\n        \n        # Drop random bodyparts from random frames\n        if rng.uniform() < 0.7:\n            mask = np.ones(len(vid), dtype=bool)\n            mask[:int(0.4 * len(mask))] = False\n            rng.shuffle(mask)\n            vid = vid[mask]\n    \n        # Set random coordinates of bodyparts to nan\n        if rng.uniform() < 0.7:\n            mask = np.ones(len(vid), dtype=bool)\n            mask[:int(0.2 * len(mask))] = False\n            rng.shuffle(mask)\n            vid.loc[:, 'x'] = np.where(mask, np.nan, vid.loc[:, 'x'])\n            rng.shuffle(mask)\n            vid.loc[:, 'y'] = np.where(mask, np.nan, vid.loc[:, 'y'])\n    \n        # Save the video\n        try:\n            os.mkdir(f\"stresstest_tracking/{lab_id}\")\n        except FileExistsError:\n            pass\n        new_path = f\"stresstest_tracking/{lab_id}/{video_id}.parquet\"\n        vid.to_parquet(new_path)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-27T10:57:12.936547Z","iopub.execute_input":"2025-09-27T10:57:12.936833Z","iopub.status.idle":"2025-09-27T10:57:12.947854Z","shell.execute_reply.started":"2025-09-27T10:57:12.936810Z","shell.execute_reply":"2025-09-27T10:57:12.946844Z"},"_kg_hide-input":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Challenge 1: Modeling for variable-sized sets of mice\n\nThe first challenge we're going to solve is the fact that we have a variable number of mice (2, 3 or 4), and that the labeled behaviors apply either to one mouse or a pair of mice.\n\nThe following function, `generate_mouse_data()`, solves this challenge. It transforms the dataset into batches. There are single-mouse batches and mouse-pair batches. Every single-mouse batch has data of only one mouse, every mouse-pair batch has data of exactly two mice. A single video frame can end up in several batches. If the frame has two visible mice, it can be part of four batches:\n- a single-mouse batch for individual behavior of mouse 1\n- a single-mouse batch for individual behavior of mouse 2\n- a mouse-pair batch for actions of mouse 1 with mouse 2 as target\n- a mouse-pair batch for actions of mouse 2 with mouse 1 as target\n\nThe features (`data`) will consist of coordinates of body parts; the metadata (`meta`) will specify which mouse is / which mice are involved.","metadata":{}},{"cell_type":"code","source":"drop_body_parts =  ['headpiece_bottombackleft', 'headpiece_bottombackright', 'headpiece_bottomfrontleft', 'headpiece_bottomfrontright', \n                    'headpiece_topbackleft', 'headpiece_topbackright', 'headpiece_topfrontleft', 'headpiece_topfrontright', \n                    'spine_1', 'spine_2',\n                    'tail_middle_1', 'tail_middle_2', 'tail_midpoint']\n\ndef generate_mouse_data(dataset, traintest, traintest_directory=None, generate_single=True, generate_pair=True):\n    \"\"\"Generate batches of data in coordinate representation.\n\n    The batches have variable length, and every batch can have other columns\n    for the labels, depending on what behaviors\n    were labeled for the batch.\n\n    Every video can produce zero, one or two batches.\n    \n    Parameters\n    ----------\n    dataset: (subset of) train.csv or test.csv dataframe\n    traintest: either 'train' or 'test'\n\n    Yields\n    ------\n    switch: either 'single' or 'pair'\n    data: dataframe containing coordinates of the body parts of a single mouse or of a pair of mice\n    meta: dataframe with columns ['video_id', 'agent_id', 'target_id', 'video_frame']\n    label: dataframe with labels (0, 1), one column per action, only if traintest == 'train'\n    actions: list of actions to be predicted for this batch, only if traintest == 'test'\n    \"\"\"\n    assert traintest in ['train', 'test']\n    if traintest_directory is None:\n        traintest_directory = f\"/kaggle/input/MABe-mouse-behavior-detection/{traintest}_tracking\"\n    for _, row in dataset.iterrows():\n        \n        # Load the video and pivot it sn that one frame = one row\n        lab_id = row.lab_id\n        if lab_id.startswith('MABe22'): continue\n        video_id = row.video_id\n\n        if type(row.behaviors_labeled) != str:\n            # We cannot use videos without labeled behaviors\n            print('No labeled behaviors:', lab_id, video_id)\n            continue\n\n        path = f\"{traintest_directory}/{lab_id}/{video_id}.parquet\"\n        vid = pd.read_parquet(path)\n        if len(np.unique(vid.bodypart)) > 5:\n            vid = vid.query(\"~ bodypart.isin(@drop_body_parts)\")\n        pvid = vid.pivot(columns=['mouse_id', 'bodypart'], index='video_frame', values=['x', 'y'])\n        if pvid.isna().any().any():\n            if verbose and traintest == 'test': print('video with missing values', video_id, traintest, len(vid), 'frames')\n        else:\n            if verbose and traintest == 'test': print('video with all values', video_id, traintest, len(vid), 'frames')\n        del vid\n        pvid = pvid.reorder_levels([1, 2, 0], axis=1).T.sort_index().T # mouse_id, body_part, xy\n        pvid /= row.pix_per_cm_approx # convert to cm\n\n        # Determine the behaviors of this video\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        # Load the annotations\n        if traintest == 'train':\n            try:\n                annot = pd.read_parquet(path.replace('train_tracking', 'train_annotation'))\n            except FileNotFoundError:\n                # MABe22 and one more training file lack annotations.\n                # We simply drop these videos.\n                continue\n\n        # Create the single_mouse dataframes: single_mouse, single_mouse_label and single_mouse_meta\n        if generate_single:\n            vid_behaviors_subset = vid_behaviors.query(\"target == 'self'\") # single-mouse behaviors of this video\n            for mouse_id_str in np.unique(vid_behaviors_subset.agent):\n                try:\n                    mouse_id = int(mouse_id_str[-1])\n                    vid_agent_actions = np.unique(vid_behaviors_subset.query(\"agent == @mouse_id_str\").action)\n                    single_mouse = pvid.loc[:, mouse_id]\n                    assert len(single_mouse) == len(pvid)\n                    single_mouse_meta = pd.DataFrame({\n                        'video_id': video_id,\n                        'agent_id': mouse_id_str,\n                        'target_id': 'self',\n                        'video_frame': single_mouse.index\n                    })\n                    if traintest == 'train':\n                        single_mouse_label = pd.DataFrame(0.0, columns=vid_agent_actions, index=single_mouse.index)\n                        annot_subset = annot.query(\"(agent_id == @mouse_id) & (target_id == @mouse_id)\")\n                        for i in range(len(annot_subset)):\n                            annot_row = annot_subset.iloc[i]\n                            single_mouse_label.loc[annot_row['start_frame']:annot_row['stop_frame'], annot_row.action] = 1.0\n                        yield 'single', single_mouse, single_mouse_meta, single_mouse_label\n                    else:\n                        if verbose: print('- test single', video_id, mouse_id)\n                        yield 'single', single_mouse, single_mouse_meta, vid_agent_actions\n                except KeyError:\n                    pass # If there is no data for the selected agent mouse, we skip the mouse.\n\n        # Create the mouse_pair dataframes: mouse_pair, mouse_label and mouse_meta\n        if generate_pair:\n            vid_behaviors_subset = vid_behaviors.query(\"target != 'self'\")\n            if len(vid_behaviors_subset) > 0:\n                for agent, target in itertools.permutations(np.unique(pvid.columns.get_level_values('mouse_id')), 2): # int8\n                    agent_str = f\"mouse{agent}\"\n                    target_str = f\"mouse{target}\"\n                    vid_agent_actions = np.unique(vid_behaviors_subset.query(\"(agent == @agent_str) & (target == @target_str)\").action)\n                    mouse_pair = pd.concat([pvid[agent], pvid[target]], axis=1, keys=['A', 'B'])\n                    assert len(mouse_pair) == len(pvid)\n                    mouse_pair_meta = pd.DataFrame({\n                        'video_id': video_id,\n                        'agent_id': agent_str,\n                        'target_id': target_str,\n                        'video_frame': mouse_pair.index\n                    })\n                    if traintest == 'train':\n                        mouse_pair_label = pd.DataFrame(0.0, columns=vid_agent_actions, index=mouse_pair.index)\n                        annot_subset = annot.query(\"(agent_id == @agent) & (target_id == @target)\")\n                        for i in range(len(annot_subset)):\n                            annot_row = annot_subset.iloc[i]\n                            mouse_pair_label.loc[annot_row['start_frame']:annot_row['stop_frame'], annot_row.action] = 1.0\n                        yield 'pair', mouse_pair, mouse_pair_meta, mouse_pair_label\n                    else:\n                        if verbose: print('- test pair', video_id, agent, target)\n                        yield 'pair', mouse_pair, mouse_pair_meta, vid_agent_actions\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-27T10:57:12.948876Z","iopub.execute_input":"2025-09-27T10:57:12.949245Z","iopub.status.idle":"2025-09-27T10:57:12.973339Z","shell.execute_reply.started":"2025-09-27T10:57:12.949211Z","shell.execute_reply":"2025-09-27T10:57:12.972395Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Challenge 2: Multiclass prediction with missing labels\n\nThis competition is a multi-class classification task. For every video_id/video_frame/agent/target combination, we may predict at most one of several actions. Every action is a class, and 'no-action' is an additional class.\n\nWe cannot use a standard multi-class estimator from scikit-learn because many values in the labels of our dataset are missing. For this reason, we train a binary classifier for every action, omitting the samples for which the target is unknown. Every binary classificator predicts a probability, and for the multiclass prediction we predict the class with the highest binary probability, if this probability is above a threshold; otherwise, we predict no action.\n","metadata":{}},{"cell_type":"code","source":"# Make the multi-class prediction\ndef predict_multiclass(pred, meta):\n    \"\"\"Derive multiclass predictions from a set of binary predictions.\n    \n    Parameters\n    pred: dataframe of predicted binary probabilities, shape (n_samples, n_actions), index doesn't matter\n    meta: dataframe with columns ['video_id', 'agent_id', 'target_id', 'video_frame'], index doesn't matter\n    \"\"\"\n    # Find the most probable class, but keep it only if its probability is above the threshold\n    ama = np.argmax(pred, axis=1)\n    ama = np.where(pred.max(axis=1) >= threshold, ama, -1)\n    ama = pd.Series(ama, index=meta.video_frame)\n    # Keep only start and stop frames\n    changes_mask = (ama != ama.shift(1)).values\n    ama_changes = ama[changes_mask]\n    meta_changes = meta[changes_mask]\n    # mask selects the start frames\n    mask = ama_changes.values >= 0 # start of action\n    mask[-1] = False\n    submission_part = pd.DataFrame({\n        'video_id': meta_changes['video_id'][mask].values,\n        'agent_id': meta_changes['agent_id'][mask].values,\n        'target_id': meta_changes['target_id'][mask].values,\n        'action': pred.columns[ama_changes[mask].values],\n        'start_frame': ama_changes.index[mask],\n        'stop_frame': ama_changes.index[1:][mask[:-1]]\n    })\n    stop_video_id = meta_changes['video_id'][1:][mask[:-1]].values\n    stop_agent_id = meta_changes['agent_id'][1:][mask[:-1]].values\n    stop_target_id = meta_changes['target_id'][1:][mask[:-1]].values\n    for i in range(len(submission_part)):\n        video_id = submission_part.video_id.iloc[i]\n        agent_id = submission_part.agent_id.iloc[i]\n        target_id = submission_part.target_id.iloc[i]\n        if stop_video_id[i] != video_id or stop_agent_id[i] != agent_id or stop_target_id[i] != target_id:\n            new_stop_frame = meta.query(\"(video_id == @video_id)\").video_frame.max() + 1\n            submission_part.iat[i, submission_part.columns.get_loc('stop_frame')] = new_stop_frame\n    assert (submission_part.stop_frame > submission_part.start_frame).all(), 'stop <= start'\n    if verbose: print('  actions found:', len(submission_part))\n    return submission_part\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-27T10:57:12.976191Z","iopub.execute_input":"2025-09-27T10:57:12.976458Z","iopub.status.idle":"2025-09-27T10:57:12.996934Z","shell.execute_reply.started":"2025-09-27T10:57:12.976437Z","shell.execute_reply":"2025-09-27T10:57:12.995833Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Challenge 3: Transforming coordinates to an invariant representation\n\nThe body part of the mice are given in cartesian coordinates. If the mice show some behavior at varying positions and with varying spatial orientation, cartesian coordinates are an inadequate representation. Our feature engineering transforms the coordinates to distances between body parts. Distances are invariant under translation and rotation.\n\nFor a single mouse, the distances indicate whether and how much it turns its head, shoulders, hip and tail left or right. For a pair of mice, the distances indicate how far the head of the first mouse is near what part of the second one, and what body parts either mouse turns towards or away from the other one.","metadata":{}},{"cell_type":"code","source":"def transform_single(single_mouse, body_parts_tracked):\n    \"\"\"Transform from cartesian coordinates to distance representation.\n\n    Parameters:\n    single_mouse: dataframe with coordinates of the body parts of one mouse\n                  shape (n_samples, n_body_parts * 2)\n                  two-level MultiIndex on columns\n    body_parts_tracked: list of body parts\n    \"\"\"\n    available_body_parts = single_mouse.columns.get_level_values(0)\n    X = pd.DataFrame({\n            f\"{part1}+{part2}\": np.square(single_mouse[part1] - single_mouse[part2]).sum(axis=1, skipna=False)\n            for part1, part2 in itertools.combinations(body_parts_tracked, 2) if part1 in available_body_parts and part2 in available_body_parts\n        })\n    X = X.reindex(columns=[f\"{part1}+{part2}\" for part1, part2 in itertools.combinations(body_parts_tracked, 2)], copy=False)\n\n    if 'ear_left' in single_mouse.columns and 'ear_right' in single_mouse.columns and 'tail_base' in single_mouse.columns:\n        shifted = single_mouse[['ear_left', 'ear_right', 'tail_base']].shift(10)\n        X = pd.concat([\n            X, \n            pd.DataFrame({\n                'speed_left': np.square(single_mouse['ear_left'] - shifted['ear_left']).sum(axis=1, skipna=False),\n                'speed_right': np.square(single_mouse['ear_right'] - shifted['ear_right']).sum(axis=1, skipna=False),\n                'speed_left2': np.square(single_mouse['ear_left'] - shifted['tail_base']).sum(axis=1, skipna=False),\n                'speed_right2': np.square(single_mouse['ear_right'] - shifted['tail_base']).sum(axis=1, skipna=False),\n            })\n        ], axis=1)\n    return X\n\ndef transform_pair(mouse_pair, body_parts_tracked):\n    \"\"\"Transform from cartesian coordinates to distance representation.\n\n    Parameters:\n    mouse_pair: dataframe with coordinates of the body parts of two mice\n                  shape (n_samples, 2 * n_body_parts * 2)\n                  three-level MultiIndex on columns\n    body_parts_tracked: list of body parts\n    \"\"\"\n    # drop_body_parts =  ['ear_left', 'ear_right',\n    #                     'headpiece_bottombackleft', 'headpiece_bottombackright', 'headpiece_bottomfrontleft', 'headpiece_bottomfrontright', \n    #                     'headpiece_topbackleft', 'headpiece_topbackright', 'headpiece_topfrontleft', 'headpiece_topfrontright', \n    #                     'tail_midpoint']\n    # if len(body_parts_tracked) > 5:\n    #     body_parts_tracked = [b for b in body_parts_tracked if b not in drop_body_parts]\n    available_body_parts_A = mouse_pair['A'].columns.get_level_values(0)\n    available_body_parts_B = mouse_pair['B'].columns.get_level_values(0)\n    X = pd.DataFrame({\n            f\"12+{part1}+{part2}\": np.square(mouse_pair['A'][part1] - mouse_pair['B'][part2]).sum(axis=1, skipna=False)\n            for part1, part2 in itertools.product(body_parts_tracked, repeat=2) if part1 in available_body_parts_A and part2 in available_body_parts_B\n        })\n    X = X.reindex(columns=[f\"12+{part1}+{part2}\" for part1, part2 in itertools.product(body_parts_tracked, repeat=2)], copy=False)\n\n    if ('A', 'ear_left') in mouse_pair.columns and ('B', 'ear_left') in mouse_pair.columns:\n        shifted_A = mouse_pair['A']['ear_left'].shift(10)\n        shifted_B = mouse_pair['B']['ear_left'].shift(10)\n        X = pd.concat([\n            X,\n            pd.DataFrame({\n                'speed_left_A': np.square(mouse_pair['A']['ear_left'] - shifted_A).sum(axis=1, skipna=False),\n                'speed_left_AB': np.square(mouse_pair['A']['ear_left'] - shifted_B).sum(axis=1, skipna=False),\n                'speed_left_B': np.square(mouse_pair['B']['ear_left'] - shifted_B).sum(axis=1, skipna=False),\n            })\n        ], axis=1)\n    return X\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-27T10:57:12.997918Z","iopub.execute_input":"2025-09-27T10:57:12.998277Z","iopub.status.idle":"2025-09-27T10:57:13.023795Z","shell.execute_reply.started":"2025-09-27T10:57:12.998223Z","shell.execute_reply":"2025-09-27T10:57:13.022830Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Cross-validation\n\nWe're now almost ready to cross-validate our models. \n\nThe following function gets as input\n- a binary classification model\n- a 2d array of features (i.e., distances between body parts); after we have dealt with variable-sized mouse sets (challenge 1) and variable-sized bodyparts sets (challenge 5), this array is rectangular.\n- a 2d array of binary labels, some elements of which may be missing\n- a 2d array of metadata so that we can match the predictions with the original video_id, agent, target and video_frame\n\nIt first computes out-of-fold predictions with a set of binary classifiers and then transforms these binary predictions into a multiclass prediction (see above).","metadata":{}},{"cell_type":"code","source":"threshold = 0.27\n\ndef cross_validate_classifier(binary_classifier, X, label, meta):\n    \"\"\"Cross-validate a binary classifier per action and a multi-class classifier over all actions.\n\n    Parameters\n    ----------\n    binary_classifier: classifier with predict_proba\n    X: 2d array-like (distance representation) of shape (n_samples, n_features)\n    label: dataframe with binary targets (one column per action, may have missing values), index doesn't matter\n    meta: dataframe with columns ['video_id', 'agent_id', 'target_id', 'video_frame'], index doesn't matter\n\n    Output\n    ------\n    appends to f1_list (binary) and submission_list (multi-class)\n    \n    \"\"\"\n    # Cross-validate a binary classifier for every action\n    oof = pd.DataFrame(index=meta.video_frame) # will get a column per action\n    for action in label.columns:\n        # Filter for samples (video frames) with a defined target (i.e., target is not nan)\n        action_mask = ~ label[action].isna().values\n        X_action = X[action_mask]\n        y_action = label[action][action_mask].values.astype(int)\n        p = y_action.mean()\n        baseline_score = p / (1 + p)\n        groups_action = meta.video_id[action_mask] # ensure validation has unseen videos\n        if len(np.unique(groups_action)) < 5:\n            continue # GroupKFold would fail with fewer than n_splits groups\n\n        if not (y_action == 0).all():\n            with warnings.catch_warnings():\n                warnings.filterwarnings('ignore', category=RuntimeWarning)\n                # Number of classes in training fold (1) does not match total number of classes (2)\n                oof_action = cross_val_predict(binary_classifier, X_action, y_action, groups=groups_action, cv=GroupKFold(), method='predict_proba')\n            oof_action = oof_action[:, 1]\n        else:\n            oof_action = np.zeros(len(y_action))\n        f1 = f1_score(y_action, (oof_action >= threshold), zero_division=0)\n        ch = '>' if f1 > baseline_score else '=' if f1 == baseline_score else '<'\n        print(f\"  F1: {f1:.3f} {ch} ({baseline_score:.3f}) {action}\")\n        f1_list.append((body_parts_tracked_str, action, f1))\n        oof_column = np.zeros(len(label))\n        oof_column[action_mask] = oof_action\n        oof[action] = oof_column\n\n    # Make the multi-class prediction\n    submission_part = predict_multiclass(oof, meta)\n    submission_list.append(submission_part)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-27T10:57:13.024987Z","iopub.execute_input":"2025-09-27T10:57:13.025270Z","iopub.status.idle":"2025-09-27T10:57:13.045256Z","shell.execute_reply.started":"2025-09-27T10:57:13.025247Z","shell.execute_reply":"2025-09-27T10:57:13.044219Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Challenge 4: A dataset that doesn't fit into memory\n\nThe competition dataset doesn't fit into memory as whole. The problem is exacerbated if we compute lots of distance in feature engineering. We tackle this challenge with the following measures:\n- Training on a subset of the data: The training dataset is highly redundant. In videos taken with 30 frames per second, the difference from one frame to the next is small. We can well afford to subsample the training data.\n- Processing the test data in batches: There is no need to have the full test dataset in memory at any time. (This decision has the drawback that the test data are read from disk several times.)\n- It helps that we split all data by body_parts_tracked (see challenge 5 below). This way, we don't even need to have the full training dataset in memory.","metadata":{}},{"cell_type":"code","source":"def submit(body_parts_tracked_str, switch_tr, binary_classifier, X_tr, label, meta):\n    \"\"\"Produce a submission file for the selected subset of the test data.\n\n    Parameters\n    ----------\n    body_parts_tracked_str: subset of body parts for filtering the test set\n    switch_tr: 'single' or 'pair'\n    binary_classifier: classifier with predict_proba\n    X_tr: training features as 2d array-like of shape (n_samples, n_features)\n    label: dataframe with binary targets (one column per action, may have missing values), index doesn't matter\n    meta: dataframe with columns ['video_id', 'agent_id', 'target_id', 'video_frame'], index doesn't matter\n\n    Output\n    ------\n    appends to submission_list\n    \n    \"\"\"\n    # Fit a binary classifier for every action\n    model_list = [] # will get a model per action\n    for action in label.columns:\n        # Filter for samples (video frames) with a defined target (i.e., target is not nan)\n        action_mask = ~ label[action].isna().values\n        y_action = label[action][action_mask].values.astype(int)\n\n        if not (y_action == 0).all():\n            model = clone(binary_classifier)\n            model.fit(X_tr[action_mask], y_action)\n            assert len(model.classes_) == 2\n            model_list.append((action, model))\n\n    # Compute test predictions in batches\n    body_parts_tracked = json.loads(body_parts_tracked_str)\n    if len(body_parts_tracked) > 5:\n        body_parts_tracked = [b for b in body_parts_tracked if b not in drop_body_parts]\n    if validate_or_submit == 'submit':\n        test_subset = test[test.body_parts_tracked == body_parts_tracked_str]\n        generator = generate_mouse_data(test_subset, 'test',\n                                        generate_single=(switch_tr == 'single'), \n                                        generate_pair=(switch_tr == 'pair'))\n    else:\n        test_subset = stresstest.query(\"body_parts_tracked == @body_parts_tracked_str\")\n        generator = generate_mouse_data(test_subset, 'test',\n                                        traintest_directory='stresstest_tracking',\n                                        generate_single=(switch_tr == 'single'),\n                                        generate_pair=(switch_tr == 'pair'))\n    if verbose: print(f\"n_videos: {len(test_subset)}\")\n    for switch_te, data_te, meta_te, actions_te in generator:\n        assert switch_te == switch_tr\n        try:\n            # Transform from coordinate representation into distance representation\n            if switch_te == 'single':\n                X_te = transform_single(data_te, body_parts_tracked) # may raise KeyError\n            else:\n                X_te = transform_pair(data_te, body_parts_tracked) # may raise KeyError\n            if verbose and len(X_te) == 0: print(\"ERROR: X_te is empty\")\n            del data_te\n    \n            # Compute binary predictions\n            pred = pd.DataFrame(index=meta_te.video_frame) # will get a column per action\n            for action, model in model_list:\n                if action in actions_te:\n                    pred[action] = model.predict_proba(X_te)[:, 1]\n            del X_te\n            # Compute multiclass predictions\n            if pred.shape[1] != 0:\n                submission_part = predict_multiclass(pred, meta_te)\n                submission_list.append(submission_part)\n            else: # this happens if there was no useful training data for the test actions\n                if verbose: print(f\"  ERROR: no useful training data\")\n        except KeyError:\n            if verbose: print(f'  ERROR: KeyError because of missing bodypart ({switch_tr})')\n            del data_te\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-27T10:57:13.046340Z","iopub.execute_input":"2025-09-27T10:57:13.046642Z","iopub.status.idle":"2025-09-27T10:57:13.072132Z","shell.execute_reply.started":"2025-09-27T10:57:13.046597Z","shell.execute_reply":"2025-09-27T10:57:13.071085Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Challenge 5: Modeling for variable sets of body parts\n\nDifferent labs have tracked different sets of body parts, but a machine learning model expects to see the same features for every sample. We solve this challenge by the principle of divide and conquer: For every set of body parts, we fit separate models.","metadata":{}},{"cell_type":"code","source":"# %%time\nf1_list = []\nsubmission_list = []\nfor section in range(1, len(body_parts_tracked_list)): # skip index 0 (MABe22)\n    body_parts_tracked_str = body_parts_tracked_list[section]\n    try:\n        body_parts_tracked = json.loads(body_parts_tracked_str)\n        print(f\"{section}. Processing videos with {body_parts_tracked}\")\n        if len(body_parts_tracked) > 5:\n            body_parts_tracked = [b for b in body_parts_tracked if b not in drop_body_parts]\n    \n        # We read all training data which match the body parts tracked\n        train_subset = train[train.body_parts_tracked == body_parts_tracked_str]\n        single_mouse_list = []\n        single_mouse_label_list = []\n        single_mouse_meta_list = []\n        mouse_pair_list = []\n        mouse_pair_label_list = []\n        mouse_pair_meta_list = []\n    \n        for switch, data, meta, label in generate_mouse_data(train_subset, 'train'):\n            if switch == 'single':\n                single_mouse_list.append(data)\n                single_mouse_meta_list.append(meta)\n                single_mouse_label_list.append(label)\n            else:\n                mouse_pair_list.append(data)\n                mouse_pair_meta_list.append(meta)\n                mouse_pair_label_list.append(label)\n    \n        # Construct a binary classifier\n        binary_classifier = make_pipeline(\n            SimpleImputer(),\n            TrainOnSubsetClassifier(\n                lightgbm.LGBMClassifier(\n                    n_estimators=100,\n                    learning_rate=0.03,\n                    min_child_samples=40,\n                    # early_stopping_round=10, \n                    verbose=-1),\n                100000)\n        )\n    \n        # Predict single-mouse actions\n        if len(single_mouse_list) > 0:\n            # Concatenate all batches\n            # The concatenation will generate label dataframes with missing values.\n            single_mouse = pd.concat(single_mouse_list)\n            single_mouse_label = pd.concat(single_mouse_label_list)\n            single_mouse_meta = pd.concat(single_mouse_meta_list)\n            del single_mouse_list, single_mouse_label_list, single_mouse_meta_list\n            assert len(single_mouse) == len(single_mouse_label)\n            assert len(single_mouse) == len(single_mouse_meta)\n            \n            # Transform the coordinate representation into a distance representation for single_mouse\n            X_tr = transform_single(single_mouse, body_parts_tracked)\n            del single_mouse\n            print(f\"{X_tr.shape=}\")\n    \n            if validate_or_submit == 'validate':\n                cross_validate_classifier(binary_classifier, X_tr, single_mouse_label, single_mouse_meta)\n            else:\n                submit(body_parts_tracked_str, 'single', binary_classifier, X_tr, single_mouse_label, single_mouse_meta)\n            del X_tr\n                \n        # Predict mouse-pair actions\n        if len(mouse_pair_list) > 0:\n            # Concatenate all batches\n            # The concatenation will generate label dataframes with missing values.\n            mouse_pair = pd.concat(mouse_pair_list)\n            mouse_pair_label = pd.concat(mouse_pair_label_list)\n            mouse_pair_meta = pd.concat(mouse_pair_meta_list)\n            del mouse_pair_list, mouse_pair_label_list, mouse_pair_meta_list\n            assert len(mouse_pair) == len(mouse_pair_label)\n            assert len(mouse_pair) == len(mouse_pair_meta)\n        \n            # Transform the coordinate representation into a distance representation for mouse_pair\n            # Use a subset of body_parts_tracked to conserve memory\n            X_tr = transform_pair(mouse_pair, body_parts_tracked)\n            del mouse_pair\n            print(f\"{X_tr.shape=}\")\n    \n            if validate_or_submit == 'validate':\n                cross_validate_classifier(binary_classifier, X_tr, mouse_pair_label, mouse_pair_meta)\n            else:\n                submit(body_parts_tracked_str, 'pair', binary_classifier, X_tr, mouse_pair_label, mouse_pair_meta)\n            del X_tr\n                \n    except Exception as e:\n        print(f'***Exception*** {e}')\n    print()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-27T10:57:13.073075Z","iopub.execute_input":"2025-09-27T10:57:13.073408Z","iopub.status.idle":"2025-09-27T11:51:56.849574Z","shell.execute_reply.started":"2025-09-27T10:57:13.073363Z","shell.execute_reply":"2025-09-27T11:51:56.848116Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"The code above probably contains bugs, but we don't want the submission to fail because of these bugs. The function `robustify` modifies the submission dataframe so that it conforms to the rules of the competition.","metadata":{}},{"cell_type":"code","source":"def robustify(submission, dataset, traintest, traintest_directory=None):\n    \"\"\"Ensure that the submission conforms to the three rules\"\"\"\n    if traintest_directory is None:\n        traintest_directory = f\"/kaggle/input/MABe-mouse-behavior-detection/{traintest}_tracking\"\n\n    # Rule 1: Ensure that start_frame >= stop_frame\n    old_submission = submission.copy()\n    submission = submission[submission.start_frame < submission.stop_frame]\n    if len(submission) != len(old_submission):\n        print(\"ERROR: Dropped frames with start >= stop\")\n    \n    # Rule 2: Avoid multiple predictions for the same frame from one agent/target pair\n    old_submission = submission.copy()\n    group_list = []\n    for _, group in submission.groupby(['video_id', 'agent_id', 'target_id']):\n        group = group.sort_values('start_frame')\n        mask = np.ones(len(group), dtype=bool)\n        last_stop_frame = 0\n        for i, (_, row) in enumerate(group.iterrows()):\n            if row['start_frame'] < last_stop_frame:\n                mask[i] = False\n            else:\n                last_stop_frame = row['stop_frame']\n        group_list.append(group[mask])\n    submission = pd.concat(group_list)\n    if len(submission) != len(old_submission):\n        print(\"ERROR: Dropped duplicate frames\")\n\n    # Rule 3: Submit something for every video\n    # Fill missing videos as in https://www.kaggle.com/code/ambrosm/mabe-validated-baseline-without-machine-learning\n    s_list = []\n    for idx, row in dataset.iterrows():\n        lab_id = row['lab_id']\n        if lab_id.startswith('MABe22'):\n            continue\n        video_id = row['video_id']\n        if (submission.video_id == video_id).any():\n            continue\n        if type(row.behaviors_labeled) != str:\n            continue\n\n        if verbose: print(f\"Video {video_id} has no predictions.\")\n        \n        # Load video\n        path = f\"{traintest_directory}/{lab_id}/{video_id}.parquet\"\n        vid = pd.read_parquet(path)\n    \n        # Determine the behaviors of this video\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                s_list.append((video_id, agent, target, action_row['action'], batch_start, batch_stop))\n\n    if len(s_list) > 0:\n        submission = pd.concat([\n            submission,\n            pd.DataFrame(s_list, columns=['video_id', 'agent_id', 'target_id', 'action', 'start_frame', 'stop_frame'])\n        ])\n        print(\"ERROR: Filled empty videos\")\n\n    submission = submission.reset_index(drop=True)\n    return submission\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-27T12:15:42.542672Z","iopub.execute_input":"2025-09-27T12:15:42.544172Z","iopub.status.idle":"2025-09-27T12:15:42.560818Z","shell.execute_reply.started":"2025-09-27T12:15:42.544124Z","shell.execute_reply":"2025-09-27T12:15:42.559537Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if validate_or_submit == 'validate':\n    # Score the oof predictions with the competition scoring function\n    submission = pd.concat(submission_list)\n    submission_robust = robustify(submission, train, 'train')\n    print(f\"# OOF score with competition metric: {score(solution, submission_robust, ''):.4f}\")\n\n    f1_df = pd.DataFrame(f1_list, columns=['body_parts_tracked_str', 'action', 'binary F1 score'])\n    print(f\"# Average of {len(f1_df)} binary F1 scores {f1_df['binary F1 score'].mean():.4f}\")\n    # with pd.option_context('display.max_rows', 500):\n    #     display(f1_df)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-27T12:15:45.227233Z","iopub.execute_input":"2025-09-27T12:15:45.227578Z","iopub.status.idle":"2025-09-27T12:16:18.319476Z","shell.execute_reply.started":"2025-09-27T12:15:45.227552Z","shell.execute_reply":"2025-09-27T12:16:18.318272Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if validate_or_submit != 'validate':\n    if len(submission_list) > 0:\n        submission = pd.concat(submission_list)\n    else:\n        submission = pd.DataFrame(\n            dict(\n                video_id=438887472,\n                agent_id='mouse1',\n                target_id='self',\n                action='rear',\n                start_frame='278',\n                stop_frame='500'\n            ), index=[44])\n    if validate_or_submit == 'submit':\n        submission_robust = robustify(submission, test, 'test')\n    else:\n        submission_robust = robustify(submission, stresstest, 'stresstest', 'stresstest_tracking')\n    submission_robust.index.name = 'row_id'\n    submission_robust.to_csv('submission.csv')\n    !head submission.csv","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-27T11:52:23.868980Z","iopub.status.idle":"2025-09-27T11:52:23.869415Z","shell.execute_reply.started":"2025-09-27T11:52:23.869202Z","shell.execute_reply":"2025-09-27T11:52:23.869222Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}